"""
   @Author      : CT
   @Description : 年度流量选品——年度聚合指标计算（Hive 侧，只负责这一部分职责）
                  - 数据来源：dwt_flow_asin 近12个月（含当月，计算年度指标）
                             + 前12个月（13~24个月前，仅供周期性判断的峰值月同比对照）
                             + dim_asin_launchtime_info（按 site_name 过滤，取权威上架时间，
                               周期性/季节性判断都靠它，不再用 dwt_flow_asin 里的月度字段）
                             只读取计算所需的窄字段，减少内存占用
                  - 计算内容：
                      bought_month_total（近一年月销总和）
                      bought_month_1 ~ bought_month_12（12个自然月月销）
                      bought_month_q1 ~ bought_month_q4（季度月销）
                      total_appear_month（全部出现月份 int 数组）
                      bought_month_peak / peak_month_arr（月销峰值 + 峰值月数组）
                      is_periodic_flag（周期性判断，算法参考信息库 dwt_ai_asin_all.py，
                                       开售不满一年不判断，直接为0）
                      is_seasonal_flag（季节性判断：(峰值-12月均值)/12月均值 > 0.8，
                                       同样要求开售满一年才计算，不满一年直接为0）
                      bsr_seen_count_total / nsr_seen_count_total（近一年BSR/NSR上榜天数总和）
                  - 结果写入 Hive dwt_flow_asin_year 表，按 site_name + date_info 分区
                    （表需人工预先建好，DDL 见文件末尾注释；每月一个分区，永久保留历史，
                    不做滚动清理——跟 dwt_flow_asin 本身的分区策略一致）
                  Doris 侧的中间表维护（dwt.{site}_flow_asin_365day）/ 年度指标同步
                  （dwt.{site}_flow_asin_365day_extra）/ 关联写入 selection 层，
                  全部由 doris_handle/dwt_flow_asin_year.py 负责，本脚本不碰 Doris
                  支持 us / uk / de 三站点
   执行示例:
       spark-submit dwt_flow_asin_year.py us 2026-05
"""
import os
import sys
import calendar
from datetime import datetime, timedelta
from functools import reduce

sys.path.append(os.path.dirname(sys.path[0]))

from pyspark.sql import functions as F
from pyspark.sql.types import IntegerType, BooleanType
from pyspark.storagelevel import StorageLevel

from utils.spark_util import SparkUtil
from utils.common_util import CommonUtil
from utils.hdfs_utils import HdfsUtils


SUPPORTED_SITES = ("us", "uk", "de")

HIVE_TABLE = "dwt_flow_asin_year"


class DwtFlowAsinYear(object):
    """
    年度流量选品聚合指标计算（Spark 侧计算 + 写入 Hive dwt_flow_asin_year）
    """

    def __init__(self, site_name="us", date_type="month", date_info="2026-05"):
        self.site_name = site_name
        self.date_type = date_type
        self.date_info = date_info

        app_name = f"{self.__class__.__name__}:{site_name}:{date_type}:{date_info}"
        self.spark = SparkUtil.get_spark_session(app_name)

        # 近12个月（含当月），用于年度指标计算
        self.last_12_month = [CommonUtil.get_month_offset(date_info, -i) for i in range(0, 12)]
        # 前12个月（13~24个月前），仅用于周期性判断的峰值月同比对照
        self.before_12_month = [CommonUtil.get_month_offset(date_info, -i) for i in range(12, 24)]

        # 上架时间基准日：当月最后一天 - 360天，早于该日期视为上架不满一年（非周期性）
        year_int, month_int = (int(x) for x in date_info.split('-'))
        last_day = calendar.monthrange(year_int, month_int)[1]
        month_end_date = datetime(year_int, month_int, last_day)
        self.launch_time_base_date = (month_end_date - timedelta(days=360)).strftime('%Y-%m-%d')
        print(f"上架时间基准日（早于此视为不满一年）：{self.launch_time_base_date}")

        self.df_last_12_month   = self.spark.sql("select 1+1")
        self.df_before_12_month = self.spark.sql("select 1+1")
        self.df_launch_time     = self.spark.sql("select 1+1")
        self.df_pivot           = self.spark.sql("select 1+1")
        self.df_peak            = self.spark.sql("select 1+1")
        self.df_bsr_nsr         = self.spark.sql("select 1+1")
        self.df_year_metrics    = self.spark.sql("select 1+1")

        self.udf_is_periodic    = F.udf(self.check_month_close, IntegerType())
        self.udf_is_consecutive = F.udf(self.is_consecutive, BooleanType())

    # ========== 周期性判断算法（沿用信息库 dwt_ai_asin_all.py 的算法）==========
    @staticmethod
    def is_close_month(m1, m2):
        """判断两个月份是否相同或相近：差值 <= 1，或者跨年相邻(12和1)"""
        if m1 is None or m2 is None:
            return False
        if abs(m1 - m2) <= 1:
            return True
        if {m1, m2} == {1, 12}:
            return True
        return False

    @staticmethod
    def check_month_close(new_list, old_list):
        """new_list（近12月峰值月）与 old_list（前12月峰值月）中任意一对相同或相近则判定周期性"""
        if not new_list or not old_list:
            return 0
        for n in new_list:
            for o in old_list:
                if DwtFlowAsinYear.is_close_month(n, o):
                    return 1
        return 0

    @staticmethod
    def is_consecutive(month_list):
        """判断峰值月份数组是否连续（含跨年，如 [11,12,1,2]）"""
        if not month_list or len(month_list) <= 1:
            return True
        month_list = sorted(month_list)
        n = len(month_list)
        normal = all(month_list[i] + 1 == month_list[i + 1] for i in range(n - 1))
        wrap = all((month_list[i] % 12) + 1 == month_list[(i + 1) % n] for i in range(n))
        return normal or wrap

    def run(self):
        self.read_data()
        self.handle_data()
        self.save_data()

    def read_data(self):
        # 近12个月：年度指标计算所需的最小字段集（bsr/nsr上榜天数；上架时间改从
        # dim_asin_launchtime_info 读，dwt_flow_asin 的月度 launch_time 字段不再用）
        sql1 = f"""
            SELECT
                asin,
                date_info,
                asin_bought_month  AS bought_month,
                bsr_seen_count_30d AS bsr_count,
                nsr_seen_count_30d AS nsr_count
            FROM dwt_flow_asin
            WHERE site_name = '{self.site_name}'
              AND date_type = '{self.date_type}'
              AND date_info IN ({CommonUtil.list_to_insql(self.last_12_month)})
        """
        self.df_last_12_month = self.spark.sql(sql1).repartition(40, 'asin').persist(StorageLevel.DISK_ONLY)
        print(f"近12月流量选品数据：{self.df_last_12_month.count()}")

        # 前12个月：仅周期性判断峰值月同比对照，只读3列
        sql2 = f"""
            SELECT
                asin,
                date_info,
                asin_bought_month AS bought_month
            FROM dwt_flow_asin
            WHERE site_name = '{self.site_name}'
              AND date_type = '{self.date_type}'
              AND date_info IN ({CommonUtil.list_to_insql(self.before_12_month)})
        """
        self.df_before_12_month = self.spark.sql(sql2).repartition(40, 'asin').persist(StorageLevel.DISK_ONLY)
        print(f"前12月流量选品数据（周期性对照用）：{self.df_before_12_month.count()}")

        # 权威上架时间：周期性/季节性判断共用同一个来源，按 site_name 过滤
        sql3 = f"""
            SELECT
                asin,
                asin_launch_time AS launch_time
            FROM dim_asin_launchtime_info
            WHERE site_name = '{self.site_name}'
        """
        self.df_launch_time = self.spark.sql(sql3).repartition(40, 'asin').persist(StorageLevel.DISK_ONLY)
        print(f"上架时间维表数据：{self.df_launch_time.count()}")

    def handle_data(self):
        self._handle_monthly_pivot()
        self._handle_peak_and_periodic()
        self._handle_bsr_nsr()
        self._handle_merge()

    def _handle_monthly_pivot(self):
        """12个自然月月销透视 + 季度月销 + 总销量 + 全部出现月份数组
        （出现月份 = 该月在 dwt_flow_asin 有记录，跟 bought_month 是否有值/是否为0无关，
        所以额外用 count(1) 透视一份"是否有记录"，而不是复用 bought_month 判空）"""
        df_pivot = self.df_last_12_month.groupBy('asin').pivot('date_info', self.last_12_month).agg(
            F.first('bought_month').alias('bought_month'),
            F.count(F.lit(1)).alias('appear_cnt'),
        )
        for month_str in self.last_12_month:
            month_num = int(month_str.split('-')[-1])
            df_pivot = df_pivot.withColumnRenamed(
                f'{month_str}_bought_month', f'bought_month_{month_num}'
            ).withColumnRenamed(
                f'{month_str}_appear_cnt', f'appear_cnt_{month_num}'
            )

        month_cols = [f'bought_month_{m}' for m in range(1, 13)]
        bought_month_total = reduce(
            lambda a, b: a + b, [F.coalesce(F.col(c), F.lit(0)) for c in month_cols]
        )
        appear_month_arr = F.array_sort(F.expr(
            "filter(array(" +
            ",".join([f"CASE WHEN appear_cnt_{m} IS NOT NULL AND appear_cnt_{m} > 0 THEN {m} END" for m in range(1, 13)]) +
            "), x -> x is not null)"
        ))

        self.df_pivot = df_pivot.withColumn(
            'bought_month_total', bought_month_total.cast('bigint')
        ).withColumn(
            'bought_month_q1',
            F.coalesce(F.col('bought_month_1'), F.lit(0)) +
            F.coalesce(F.col('bought_month_2'), F.lit(0)) +
            F.coalesce(F.col('bought_month_3'), F.lit(0))
        ).withColumn(
            'bought_month_q2',
            F.coalesce(F.col('bought_month_4'), F.lit(0)) +
            F.coalesce(F.col('bought_month_5'), F.lit(0)) +
            F.coalesce(F.col('bought_month_6'), F.lit(0))
        ).withColumn(
            'bought_month_q3',
            F.coalesce(F.col('bought_month_7'), F.lit(0)) +
            F.coalesce(F.col('bought_month_8'), F.lit(0)) +
            F.coalesce(F.col('bought_month_9'), F.lit(0))
        ).withColumn(
            'bought_month_q4',
            F.coalesce(F.col('bought_month_10'), F.lit(0)) +
            F.coalesce(F.col('bought_month_11'), F.lit(0)) +
            F.coalesce(F.col('bought_month_12'), F.lit(0))
        ).withColumn(
            'total_appear_month', appear_month_arr
        ).persist(StorageLevel.DISK_ONLY)
        print(f"12个自然月透视+季度聚合完成：{self.df_pivot.count()}")

    def _handle_peak_and_periodic(self):
        """月销峰值 + 峰值月数组 + 周期性判断"""
        # 近12月峰值（允许多个月并列最高，全部收集）
        df_max_last = self.df_last_12_month.groupBy('asin').agg(
            F.max('bought_month').alias('bought_month_peak')
        )
        df_peak_last = self.df_last_12_month.join(df_max_last, 'asin', 'left') \
            .filter(F.col('bought_month') == F.col('bought_month_peak')) \
            .withColumn('month', F.split(F.col('date_info'), '-')[1].cast('int')) \
            .groupBy('asin', 'bought_month_peak').agg(
                F.array_sort(F.collect_set('month')).alias('peak_month_arr')
            )

        # 前12月峰值月数组（仅用于同比对照，不输出）
        df_max_before = self.df_before_12_month.groupBy('asin').agg(
            F.max('bought_month').alias('bought_month_peak_before')
        )
        df_peak_before = self.df_before_12_month.join(df_max_before, 'asin', 'left') \
            .filter(F.col('bought_month') == F.col('bought_month_peak_before')) \
            .withColumn('month', F.split(F.col('date_info'), '-')[1].cast('int')) \
            .groupBy('asin').agg(
                F.array_sort(F.collect_set('month')).alias('peak_month_arr_before')
            )

        self.df_peak = df_peak_last.join(df_peak_before, 'asin', 'left') \
            .join(self.df_launch_time, 'asin', 'left') \
            .withColumn(
                'is_periodic_flag',
                F.when(
                    F.col('launch_time') >= F.lit(self.launch_time_base_date), F.lit(0)
                ).when(
                    F.size('peak_month_arr') > 6, F.lit(0)
                ).when(
                    (F.size('peak_month_arr') > 3) & (~self.udf_is_consecutive(F.col('peak_month_arr'))), F.lit(0)
                ).otherwise(
                    self.udf_is_periodic(F.col('peak_month_arr'), F.col('peak_month_arr_before'))
                )
            ).withColumn(
                # 峰值月超过6个月视为没有明显峰值，置为null（is_periodic_flag已用原始数组判断过，不受影响）
                'peak_month_arr',
                F.when(F.size('peak_month_arr') > 6, F.lit(None)).otherwise(F.col('peak_month_arr'))
            ).select(
                'asin', 'bought_month_peak', 'peak_month_arr', 'is_periodic_flag', 'launch_time'
            ).persist(StorageLevel.DISK_ONLY)
        print(f"月销峰值+周期性判断完成：{self.df_peak.count()}")

    def _handle_bsr_nsr(self):
        """近一年 BSR/NSR 上榜天数总和"""
        self.df_bsr_nsr = self.df_last_12_month.groupBy('asin').agg(
            F.sum(F.coalesce(F.col('bsr_count'), F.lit(0))).cast('bigint').alias('bsr_seen_count_total'),
            F.sum(F.coalesce(F.col('nsr_count'), F.lit(0))).cast('bigint').alias('nsr_seen_count_total')
        ).persist(StorageLevel.DISK_ONLY)

    def _handle_merge(self):
        """合并透视指标 + 峰值/周期性 + BSR/NSR，计算季节性判断"""
        self.df_year_metrics = self.df_pivot.join(
            self.df_peak, 'asin', 'left'
        ).join(
            self.df_bsr_nsr, 'asin', 'left'
        ).withColumn(
            'bought_month_avg', F.col('bought_month_total') / F.lit(12)
        ).withColumn(
            'is_seasonal_flag',
            F.when(
                # 开售不满一年不判断季节性，跟周期性判断共用同一个 launch_time_base_date 基准
                F.col('launch_time') >= F.lit(self.launch_time_base_date), F.lit(0)
            ).when(
                F.col('bought_month_avg') > 0,
                F.when(
                    (F.col('bought_month_peak') - F.col('bought_month_avg')) / F.col('bought_month_avg') > 0.8,
                    F.lit(1)
                ).otherwise(F.lit(0))
            ).otherwise(F.lit(0))
        ).select(
            F.col('asin'),
            F.col('bought_month_total'),
            F.coalesce(F.col('bought_month_1'), F.lit(0)).cast(IntegerType()).alias('bought_month_1'),
            F.coalesce(F.col('bought_month_2'), F.lit(0)).cast(IntegerType()).alias('bought_month_2'),
            F.coalesce(F.col('bought_month_3'), F.lit(0)).cast(IntegerType()).alias('bought_month_3'),
            F.coalesce(F.col('bought_month_4'), F.lit(0)).cast(IntegerType()).alias('bought_month_4'),
            F.coalesce(F.col('bought_month_5'), F.lit(0)).cast(IntegerType()).alias('bought_month_5'),
            F.coalesce(F.col('bought_month_6'), F.lit(0)).cast(IntegerType()).alias('bought_month_6'),
            F.coalesce(F.col('bought_month_7'), F.lit(0)).cast(IntegerType()).alias('bought_month_7'),
            F.coalesce(F.col('bought_month_8'), F.lit(0)).cast(IntegerType()).alias('bought_month_8'),
            F.coalesce(F.col('bought_month_9'), F.lit(0)).cast(IntegerType()).alias('bought_month_9'),
            F.coalesce(F.col('bought_month_10'), F.lit(0)).cast(IntegerType()).alias('bought_month_10'),
            F.coalesce(F.col('bought_month_11'), F.lit(0)).cast(IntegerType()).alias('bought_month_11'),
            F.coalesce(F.col('bought_month_12'), F.lit(0)).cast(IntegerType()).alias('bought_month_12'),
            F.col('bought_month_q1').cast(IntegerType()),
            F.col('bought_month_q2').cast(IntegerType()),
            F.col('bought_month_q3').cast(IntegerType()),
            F.col('bought_month_q4').cast(IntegerType()),
            F.concat_ws(',', F.col('total_appear_month')).alias('total_appear_month'),
            F.col('bought_month_peak').cast(IntegerType()),
            # peak_month_arr 在该 asin 全年 bought_month 都为空时会是 NULL 数组，
            # concat_ws 对 NULL 数组的结果是空字符串 ''，这里用 nullif 转回 NULL
            F.expr("nullif(concat_ws(',', peak_month_arr), '')").alias('peak_month_arr'),
            F.coalesce(F.col('is_periodic_flag'), F.lit(0)).cast(IntegerType()).alias('is_periodic_flag'),
            F.col('is_seasonal_flag').cast(IntegerType()),
            F.col('bsr_seen_count_total'),
            F.col('nsr_seen_count_total'),
        ).repartition(40, 'asin').persist(StorageLevel.DISK_ONLY)
        print(f"年度聚合指标计算完成，共 {self.df_year_metrics.count()} 个ASIN")

        self.df_last_12_month.unpersist()
        self.df_before_12_month.unpersist()
        self.df_launch_time.unpersist()
        self.df_pivot.unpersist()
        self.df_peak.unpersist()
        self.df_bsr_nsr.unpersist()

    def save_data(self):
        """写入 Hive dwt_flow_asin_year 表，按 site_name + date_info 分区（表需人工预先建好）
        先删除该分区目录再 append 写入，避免重跑时分区内数据重复"""
        df_save = self.df_year_metrics.withColumn(
            'site_name', F.lit(self.site_name)
        ).withColumn(
            'date_info', F.lit(self.date_info)
        )

        hdfs_path = f"/home/{SparkUtil.DEF_USE_DB}/dwt/{HIVE_TABLE}/site_name={self.site_name}/date_info={self.date_info}"
        print(f"清除hdfs目录中.....{hdfs_path}")
        HdfsUtils.delete_hdfs_file(hdfs_path)

        df_save.write.saveAsTable(name=HIVE_TABLE, format='hive', mode='append', partitionBy=['site_name', 'date_info'])
        print(f"年度聚合指标写入 Hive {HIVE_TABLE} 完成：site_name={self.site_name}, date_info={self.date_info}")


if __name__ == "__main__":
    site_name = sys.argv[1]
    date_info = sys.argv[2]
    assert site_name in SUPPORTED_SITES, f"不支持的站点：{site_name}，仅支持 us/uk/de"

    DwtFlowAsinYear(site_name=site_name, date_type='month', date_info=date_info).run()
