"""
   @Author      : CT
   @Description : 信息库月度整合脚本
                  - 从 dwt_flow_asin 读取当月符合条件的ASIN及详情字段（含 mom/yoy）
                  - 过滤禁用分类，筛选信息库候选ASIN（月销>=50，asin_type in (0,1)）
                  - is_new_flag：anti join Doris us_ai_asin_detail 存量asin
                  - is_ascending_flag：近6月 asin_bought_mom>0 满6次
                  - 关联店铺（dwt_fb_base_report）、ODS评论（ods_asin_detail）
                  - Hive dwt_ai_asin 备份 + Doris selection.{site}_ai_asin_detail 增量写入
"""
import os
import sys

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

from pyspark.sql import functions as F
from pyspark.sql.window import Window
from pyspark.storagelevel import StorageLevel
from utils.spark_util import SparkUtil
from utils.DorisHelper import DorisHelper
from utils.common_util import CommonUtil
from utils.hdfs_utils import HdfsUtils
from utils.db_util import DBUtil


class DwtAiAsin(object):

    def __init__(self, site_name="us", date_type="month", date_info="2024-10"):
        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)

        # 近6个月（含当月）用于 is_ascending_flag
        self.last_6_month = [CommonUtil.get_month_offset(date_info, -i) for i in range(0, 6)]

        self.df_base_asin      = self.spark.sql("select 1+1")
        self.df_existing_asin  = self.spark.sql("select 1+1")
        self.df_ascending_flag = self.spark.sql("select 1+1")
        self.df_fb_info        = self.spark.sql("select 1+1")
        self.df_ods_detail     = self.spark.sql("select 1+1")
        self.df_save           = self.spark.sql("select 1+1")

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

    def read_data(self):
        # ── 1. 当月流量选品──
        sql1 = f"""
            SELECT
                asin,
                asin_weight                                     AS weight,
                asin_bought_month                               AS bought_month,
                asin_category_desc                              AS category,
                asin_img_url                                    AS img,
                asin_title                                      AS title,
                asin_brand_name                                 AS brand,
                account_name,
                asin_buy_box_seller_type                        AS buy_box_seller_type,
                asin_launch_time                                AS launch_time,
                asin_img_num                                    AS img_num,
                CASE WHEN variation_num > 0 THEN 1 ELSE 0 END  AS variation_flag,
                variation_num,
                asin_ao_val                                     AS ao_val,
                category_first_id                               AS category_id,
                category_id                                     AS category_current_id,
                parent_asin,
                first_category_rank                             AS bsr_rank,
                asin_price                                      AS price,
                asin_rating                                     AS rating,
                asin_total_comments                             AS total_comments,
                asin_launch_time_type                           AS launch_time_type,
                asin_describe                                   AS describe,
                asin_bought_mom                                 AS bought_month_mom,
                asin_bought_yoy                                 AS bought_month_yoy
            FROM dwt_flow_asin
            WHERE site_name        = '{self.site_name}'
              AND date_type        = '{self.date_type}'
              AND date_info        = '{self.date_info}'
              AND asin_type        IN (0, 1)
              AND asin_bought_month >= 50
        """
        self.df_base_asin = self.spark.sql(sql1).repartition(40, 'asin')

        # 禁用分类过滤（category_id + category 双重）
        conn_info = DBUtil.get_connection_info("mysql", "us")
        sql_cat_id = f"""
            SELECT DISTINCT category_id FROM category_full_name a
            WHERE EXISTS (
                SELECT 1 FROM category_disable_config b
                WHERE b.id_path IS NOT NULL
                  AND a.id_path LIKE CONCAT(b.id_path, '%')
                  AND a.site = b.site
            ) AND a.site = '{self.site_name}'
        """
        df_filter_cat_id = SparkUtil.read_jdbc_query(
            self.spark, conn_info["url"], conn_info["pwd"], conn_info["username"], sql_cat_id
        )

        sql_cat_desc = f"""
            SELECT DISTINCT name_path AS category FROM category_disable_config
            WHERE site = '{self.site_name}'
        """
        df_filter_cat_desc = SparkUtil.read_jdbc_query(
            self.spark, conn_info["url"], conn_info["pwd"], conn_info["username"], sql_cat_desc
        )

        self.df_base_asin = self.df_base_asin \
            .join(df_filter_cat_id,   'category_id', 'left_anti') \
            .join(df_filter_cat_desc, 'category',    'left_anti') \
            .repartition(40, 'asin') \
            .persist(StorageLevel.DISK_ONLY)
        print(f"当月信息库ASIN：{self.df_base_asin.count()}")

        # ── 2. 存量asin（Doris us_ai_asin_detail），用于 is_new_flag ──
        sql_exist = f"SELECT asin FROM `selection`.`{self.site_name}_ai_asin_detail`"
        self.df_existing_asin = DorisHelper.spark_import_with_sql(
            self.spark, sql_exist
        ).repartition(40, 'asin').persist(StorageLevel.DISK_ONLY)
        print(f"Doris 信息库存量ASIN：{self.df_existing_asin.count()}")

        # ── 3. 近6月 asin_bought_mom，用于 is_ascending_flag ──
        sql_asc = f"""
            SELECT asin, asin_bought_mom
            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_6_month)})
        """
        self.df_ascending_flag = self.spark.sql(sql_asc).repartition(40, 'asin').persist(StorageLevel.DISK_ONLY)

        # ── 4. 店铺信息 ──
        sql_fb = f"""
            SELECT account_name, seller_id, fb_country_name, business_addr AS account_addr
            FROM dwt_fb_base_report
            WHERE site_name = '{self.site_name}'
              AND date_type = '{self.date_type}'
              AND date_info = '{self.date_info}'
        """
        self.df_fb_info = self.spark.sql(sql_fb).dropDuplicates(['account_name']).persist(StorageLevel.DISK_ONLY)

        # ── 5. ODS 评论/产品详情（取 updated_at 最新一条）──
        sql_ods = f"""
            SELECT asin, review_json_list, product_json, product_detail_json, updated_at
            FROM ods_asin_detail
            WHERE site_name = '{self.site_name}'
              AND date_type = '{self.date_type}'
              AND date_info = '{self.date_info}'
        """
        window_ods = Window.partitionBy('asin').orderBy(F.col('updated_at').desc_nulls_last())
        self.df_ods_detail = self.spark.sql(sql_ods) \
            .withColumn('rn', F.row_number().over(window_ods)) \
            .filter('rn = 1').drop('rn', 'updated_at') \
            .repartition(40, 'asin').persist(StorageLevel.DISK_ONLY)

    def handle_data(self):
        # is_new_flag：不在 Doris 存量中的视为新增
        self.df_base_asin = self.df_base_asin \
            .join(self.df_existing_asin.withColumn('_exist', F.lit(0)), 'asin', 'left') \
            .withColumn('is_new_flag', F.when(F.col('_exist').isNull(), F.lit(1)).otherwise(F.lit(0))) \
            .drop('_exist')

        # is_ascending_flag：近6月 mom 全部 > 0 且出现6次
        df_asc = self.df_ascending_flag.groupBy('asin').agg(
            F.count('asin').alias('month_cnt'),
            F.sum(F.when(F.col('asin_bought_mom') > 0, 1).otherwise(0)).alias('rising_cnt')
        ).withColumn(
            'is_ascending_flag',
            F.when((F.col('month_cnt') == 6) & (F.col('rising_cnt') == 6), F.lit(1)).otherwise(F.lit(0))
        ).drop('month_cnt', 'rising_cnt')

        # 汇总关联
        self.df_save = self.df_base_asin \
            .join(df_asc,             'asin', 'left') \
            .join(self.df_fb_info,    'account_name', 'left') \
            .join(self.df_ods_detail, 'asin', 'left') \
            .fillna({'is_new_flag': 0, 'is_ascending_flag': 0})

    def save_data(self):
        # 字段对齐 us_ai_asin_detail
        self.df_save = self.df_save.select(
            F.col('asin'),
            F.col('weight'),
            F.col('bought_month'),
            F.col('category'),
            F.col('img'),
            F.col('title'),
            F.col('brand'),
            F.col('account_name'),
            F.col('account_addr'),
            F.col('buy_box_seller_type'),
            F.col('launch_time'),
            F.col('img_num'),
            F.col('variation_flag'),
            F.col('variation_num'),
            F.col('ao_val'),
            F.col('category_id'),
            F.col('category_current_id'),
            F.col('parent_asin'),
            F.col('bsr_rank'),
            F.col('price'),
            F.col('rating'),
            F.col('total_comments'),
            F.col('seller_id'),
            F.col('fb_country_name'),
            F.col('review_json_list'),
            F.col('launch_time_type'),
            F.col('describe'),
            F.col('product_json'),
            F.col('product_detail_json'),
            F.col('bought_month_mom'),
            F.col('bought_month_yoy'),
            F.col('is_new_flag'),
            F.col('is_ascending_flag'),
            F.lit(0).alias('chat_flag'),
            # 分区字段必须在最后
            F.lit(self.site_name).alias('site_name'),
            F.lit(self.date_type).alias('date_type'),
            F.lit(self.date_info).alias('date_info'),
        ).repartition(10).persist(StorageLevel.DISK_ONLY)
        print(f"最终输出ASIN数量：{self.df_save.count()}")

        # ── Hive 写入（备份，含分区字段）──
        hive_tb   = "dwt_ai_asin"
        hdfs_path = CommonUtil.build_hdfs_path(hive_tb, partition_dict={
            "site_name": self.site_name,
            "date_type": self.date_type,
            "date_info": self.date_info,
        })
        HdfsUtils.delete_file_in_folder(hdfs_path)
        print(f"写入 Hive {hive_tb}，路径：{hdfs_path}")
        self.df_save.write.saveAsTable(
            name=hive_tb, format='hive', mode='append',
            partitionBy=['site_name', 'date_type', 'date_info']
        )

        # ── Doris 写入（us_ai_asin_detail 字段，chat_flag=0 增量）──
        doris_columns = (
            "asin, site_name, weight, bought_month, category, img, title, brand, "
            "account_name, account_addr, buy_box_seller_type, launch_time, img_num, "
            "variation_flag, variation_num, ao_val, category_id, category_current_id, "
            "parent_asin, bsr_rank, price, rating, total_comments, seller_id, "
            "fb_country_name, review_json_list, launch_time_type, `describe`, "
            "product_json, product_detail_json, bought_month_mom, bought_month_yoy, "
            "is_new_flag, is_ascending_flag, chat_flag"
        )
        try:
            DorisHelper.spark_export_with_columns(
                df_save=self.df_save.drop('date_type', 'date_info'),
                db_name='selection',
                table_name=f'{self.site_name}_ai_asin_detail',
                table_columns=doris_columns,
            )
            print("success!")
            CommonUtil.send_wx_msg(['chenyuanjie', 'wujicang'], 'ASIN信息库月度数据写入Doris成功',
                                    f'详情：{self.site_name}_ai_asin_detail {self.site_name} {self.date_type} {self.date_info}')
        except Exception as e:
            print("An error occurred while writing to Doris:", str(e))
            CommonUtil.send_wx_msg(['chenyuanjie'], '⚠ ASIN信息库月度数据写入Doris失败',
                                    f'详情：{self.site_name}_ai_asin_detail {self.site_name} {self.date_type} {self.date_info}')


if __name__ == "__main__":
    site_name = sys.argv[1]
    date_type = sys.argv[2]
    date_info = sys.argv[3]
    handle_obj = DwtAiAsin(site_name=site_name, date_type=date_type, date_info=date_info)
    handle_obj.run()
