import os
import random
import sys
import time
import traceback

import pandas as pd
from pyspark.storagelevel import StorageLevel
sys.path.append(os.path.dirname(sys.path[0]))  # 上级目录
from utils.templates import Templates
# from ..utils.templates import Templates
# from AmazonSpider.pyspark_job.utils.templates import Templates
from pyspark.sql.types import StructType, StructField, IntegerType, StringType
# 分组排序的udf窗口函数
from pyspark.sql.window import Window
from pyspark.sql import functions as F
from utils.db_util import DbTypes, DBUtil
from utils.common_util import CommonUtil
from datetime import datetime, timedelta
from utils.hdfs_utils import HdfsUtils
from utils.secure_db_client import get_remote_engine
from utils.DorisHelper import DorisHelper

# 数字/虚拟类目排除列表 — 筛选条件用
EXCLUDE_CATEGORIES = (
    "audible", "books", "digital-text", "dmusic",
    "mobile-apps", "movies-tv", "music", "software", "videogames"
)

# asin_is_need 判断中需要过滤的类目（对齐 dwt.handle_asin_is_hide）
NEED_FILTER_CATEGORIES = (
    'mobile-apps', 'audible', 'books', 'music', 'dmusic', 'digital-text',
    'magazines', 'movies-tv', 'software', 'videogames',
    'amazon-devices', 'boost', 'us-live-explorations', 'amazon-renewed'
)

# asin_is_hide 中 category_id 级别的隐藏节点
HIDE_CATEGORY_IDS = (
    '21393128011', '21377129011', '21377127011',
    '21377130011', '21388218011', '21377132011'
)


class DwdNsrBsrKeepaAsin(Templates):

    def __init__(self, site_name="us", date_type="day", date_info="2026-02-02"):
        super().__init__()
        self.site_name = site_name
        self.date_type = date_type
        self.date_info = date_info
        self.db_save = f"dwd_nsr_bsr_keepa_asin"
        self.db_save_cate = f"dwd_asin_cate_flag"
        self.spark = self.create_spark_object(app_name=f"{self.db_save}: {self.site_name}, {self.date_type}, {self.date_info}")
        self.df_save = self.spark.sql(f"select 1+1;")
        self.df_save_asin_cate = self.spark.sql(f"select 1+1;")
        self.df_asin_nsr = self.spark.sql(f"select 1+1;")
        self.df_asin_bsr = self.spark.sql(f"select 1+1;")
        self.df_asin_keepa = self.spark.sql(f"select 1+1;")
        self.partitions_by = ['site_name', 'date_type', 'date_info']
        self.reset_partitions(partitions_num=5)
        self.get_date_info_tuple()
        self.engine_mysql = get_remote_engine(
            site_name=self.site_name,
            db_type='mysql'
        )

    @staticmethod
    def get_date_30_days_ago(date_str: str, date_format: str = "%Y-%m-%d") -> str:
        # 将字符串转换为日期对象
        date_obj = datetime.strptime(date_str, date_format)

        # 获取 30 天前的日期
        thirty_days_ago = date_obj - timedelta(days=30)

        # 转换为字符串并返回
        return thirty_days_ago.strftime(date_format)

    # # 测试
    # date_input = "2026-02-03"  # 示例日期
    # result = get_date_30_days_ago(date_input)
    # print(result)

    @staticmethod
    def is_saturday(date_str: str, fmt: str = "%Y-%m-%d") -> bool:
        dt = datetime.strptime(date_str, fmt)
        # weekday(): 周一=0 ... 周六=5 周日=6
        return dt.weekday() == 5

    # # 示例
    # print(is_saturday("2026-02-07"))  # True
    # print(is_saturday("2026-02-06"))  # False

    def read_data(self):
        if self.date_type == 'day':
            thirty_days_ago = self.get_date_30_days_ago(date_str=self.date_info)
        else:
            thirty_days_ago = ''
        print(f"1.1 读取最近30天的bsr的asin")
        sql_bsr = f"select asin, date_info, 1 as asin_cate_flag, category_id from dim_bsr_asin_rank_history where site_name='{self.site_name}' and date_info between '{thirty_days_ago}' and '{self.date_info}'"
        if self.date_type == 'month':
            sql_bsr = f"select asin, date_info, 1 as asin_cate_flag, category_id from dim_bsr_asin_rank_history where site_name='{self.site_name}' and date_info in {self.date_info_tuple}"
        print("sql_bsr:", sql_bsr)
        self.df_asin_bsr = self.spark.sql(sqlQuery=sql_bsr).cache()
        self.df_asin_bsr.show(10, truncate=False)
        print(f"1.2 读取最近30天的nsr的asin")
        sql_nsr = f"select asin, date_info, 2 as asin_cate_flag, category_id from dim_nsr_asin_rank_history where site_name='{self.site_name}' and date_info between '{thirty_days_ago}' and '{self.date_info}'"
        if self.date_type == 'month':
            sql_nsr = f"select asin, date_info, 2 as asin_cate_flag, category_id from dim_nsr_asin_rank_history where site_name='{self.site_name}' and date_info in {self.date_info_tuple}"
        print("sql_nsr:", sql_nsr)
        self.df_asin_nsr = self.spark.sql(sqlQuery=sql_nsr).cache()
        self.df_asin_nsr.show(10, truncate=False)
        print(f"1.3 读取最近30天的keepa的asin")
        sql_keepa = f"""
        select distinct(asin), date_info, 3 as asin_cate_flag, null as category_id from ods_keepa_finder_asin where site_name='{self.site_name}' and date_info between '{thirty_days_ago}' and '{self.date_info}' 
         union all 
         select distinct(asin), '{self.date_info}' as date_info, 4 as asin_cate_flag, null as category_id from ods_keepa_finder_task_asin  where site_name='{self.site_name}' and date_info between '{thirty_days_ago}' and '{self.date_info}' 
        """
        if self.date_type == 'month':
            sql_keepa = f"""
            select asin, date_info, 3 as asin_cate_flag, null as category_id from ods_keepa_finder_asin where site_name='{self.site_name}' and date_info in {self.date_info_tuple} 
            union all 
         select distinct(asin), date_info, 4 as asin_cate_flag, null as category_id from ods_keepa_finder_task_asin  where site_name='{self.site_name}' and date_info in {self.date_info_tuple}
            """
        print("sql_keepa:", sql_keepa)
        self.df_asin_keepa = self.spark.sql(sqlQuery=sql_keepa).cache()
        self.df_asin_keepa.show(10, truncate=False)
        print(f"2.1 读取最近bsr分类树")
        sql_bsr = f"select category_id, category_first_id, lower(en_name) as en_name from dim_bsr_category_tree WHERE site_name ='{self.site_name}';"
        self.df_bsr_tree = self.spark.sql(sqlQuery=sql_bsr).cache()
        self.df_bsr_tree.show(10, truncate=False)
        print(f"2.2 读取最近nsr分类树")
        sql_nsr = f"select category_id, category_first_id, lower(en_name) as en_name  from dim_nsr_category_tree WHERE site_name ='{self.site_name}';"
        self.df_nsr_tree = self.spark.sql(sqlQuery=sql_nsr).cache()
        self.df_nsr_tree.show(10, truncate=False)
        print(f"2.3 从mysql读取隐藏分类")
        # 流量选品模块隐藏分类：category_full_name + category_disable_config 按 id_path 前缀匹配
        # （对齐 dwt.handle_asin_is_hide）；配置表集中存在 us(selection) 库，固定用 us 连接，
        # 实际站点靠 site 字段过滤，不能用 self.engine_mysql（按 self.site_name 连到对应站点库）
        engine_mysql_us = get_remote_engine(site_name="us", db_type='mysql')
        sql_hide = f"""
            SELECT DISTINCT category_id, 1 as hide_flag FROM category_full_name a
            WHERE EXISTS (
                SELECT 1 FROM category_disable_config b
                WHERE b.site = a.site AND b.module = '流量选品:名称前缀筛选'
                  AND a.id_path LIKE CONCAT(b.id_path, '%')
            ) AND a.site = '{self.site_name}'
        """
        pdf_hide_cate = engine_mysql_us.read_sql(sql_hide)
        if pdf_hide_cate.shape[0]:
            schema = StructType([
                StructField('category_id', StringType(), True),
                StructField('hide_flag', IntegerType(), True),
            ])
            self.df_hide_cate = self.spark.createDataFrame(pdf_hide_cate, schema=schema).cache()
            print(f"self.df_hide_cate: {self.df_hide_cate.count()}")
            self.df_hide_cate.show(10, truncate=False)

    def handle_data_hide_category(self):
        df_bsr_cate = self.df_asin_bsr.select("asin", "asin_cate_flag", "category_id", "date_info")
        df_bsr_cate = df_bsr_cate.join(self.df_bsr_tree, on='category_id', how='left').join(self.df_hide_cate, on='category_id', how='left')
        df_bsr_cate.show(10, truncate=False)

        df_nsr_cate = self.df_asin_nsr.select("asin", "asin_cate_flag", "category_id", "date_info")
        df_nsr_cate = df_nsr_cate.join(self.df_nsr_tree, on='category_id', how='left').join(self.df_hide_cate, on='category_id', how='left')
        df_nsr_cate.show(10, truncate=False)

        df_keepa_cate = self.df_asin_keepa.select("asin", "asin_cate_flag", "category_id", "date_info")

        df_union_cate = df_bsr_cate.unionByName(df_nsr_cate).unionByName(df_keepa_cate, allowMissingColumns=True)
        df_union_cate.show(10, truncate=False)
        print(f"df_union_cate-去重前: {df_union_cate.count()}")
        window = Window.partitionBy(['asin']).orderBy(F.asc_nulls_last("asin_cate_flag"))
        df_union_cate = df_union_cate.withColumn(
            "rk", F.row_number().over(window=window)
        ).filter("rk=1").drop("rk").cache()
        df_union_cate.show(10, truncate=False)
        print(f"df_union_cate-去重后: {df_union_cate.count()}")

        df_union_cate = df_union_cate.withColumn(
            "hide_first",
            F.when(F.col("category_first_id").isin(*NEED_FILTER_CATEGORIES), 1).otherwise(0)
        )
        df_union_cate = df_union_cate.withColumn(
            "hide_flag",
            F.coalesce(F.col("hide_flag"), F.lit(0))
        )
        # 数字虚拟类目(hide_first，对齐dwt的asin_is_need)优先判为2，隐藏分类配置(hide_flag，对齐asin_is_hide)判为3
        df_union_cate = df_union_cate.withColumn(
            "asin_type",
            F.when(F.col("hide_first") == 1, 2).when(F.col("hide_flag") == 1, 3).otherwise(0)
        )
        df_union_cate.show(10, truncate=False)
        df_union_cate.groupby(["asin_type"]).agg(
            F.count("asin")
        ).show(10)
        print(f"最终满足要导出的数量为： {df_union_cate.filter('asin_type=0').count()}")
        df_union_cate = df_union_cate.select("asin", "category_first_id", "asin_type")
        return df_union_cate

    def handle_data(self):
        df_union_cate = self.handle_data_hide_category()
        df_bsr = self.df_asin_bsr.select("asin", "asin_cate_flag")
        df_nsr = self.df_asin_nsr.select("asin", "asin_cate_flag")
        df_keepa = self.df_asin_keepa.select("asin", "asin_cate_flag")
        # df_keepa = self.df_asin_keepa.select("asin", "asin_cate_flag")
        df_union = df_bsr.unionByName(df_nsr).unionByName(df_keepa)
        # df_union = df_bsr.unionByName(df_nsr)
        df_result = (
            df_union
                .groupBy("asin")
                .agg(
                F.sort_array(F.collect_set("asin_cate_flag")).alias("flag_list")
            )
                .withColumn(
                "asin_cate_flag",
                F.concat_ws(",", F.col("flag_list"))
            )
                .drop("flag_list")
        )
        df_result.show(10, truncate=False)
        df_result.groupBy("asin_cate_flag").count().orderBy("count", ascending=False).show(truncate=False)
        self.df_save_asin_cate = df_result
        df_bsr_agg = self.df_asin_bsr.groupBy(['asin']).agg(
            F.max("date_info").alias("bsr_latest_date"),
            # F.count("date_info").alias("bsr_30day_count"),
            F.countDistinct("date_info").alias("bsr_30day_count"),
        )
        df_nsr_agg = self.df_asin_nsr.groupBy(['asin']).agg(
            F.max("date_info").alias("nsr_latest_date"),
            # F.count("date_info").alias("nsr_30day_count"),
            F.countDistinct("date_info").alias("nsr_30day_count"),
        )
        self.df_save_asin_cate = self.df_save_asin_cate.join(
            df_bsr_agg, on=['asin'], how='left'
        ).join(
            df_nsr_agg, on=['asin'], how='left'
        )
        self.df_save_asin_cate = self.df_save_asin_cate.withColumn("site_name", F.lit(self.site_name))
        self.df_save_asin_cate = self.df_save_asin_cate.withColumn("date_type", F.lit(self.date_type))
        self.df_save_asin_cate = self.df_save_asin_cate.withColumn("date_info", F.lit(self.date_info))
        self.df_save_asin_cate.show(10, truncate=False)
        self.save_data_common(
            df_save=self.df_save_asin_cate,
            db_save=self.db_save_cate,
            partitions_num=self.partitions_num,
            partitions_by=self.partitions_by,
            # partitions_by=["site_name"],
        )
        # self.df_save_asin_cate = self.df_save_asin_cate.withColumn("date_type", F.lit(self.date_type))
        # self.df_save_asin_cate = self.df_save_asin_cate.withColumn("date_info", F.lit(self.date_info))

        # df_asin_today = self.df_save.filter(f"date_info='{self.date_info}'")
        is_saturday_flag = True if self.date_type =='month' else self.is_saturday(self.date_info)
        self.df_save = self.df_asin_bsr.unionByName(self.df_asin_nsr, allowMissingColumns=True).unionByName(self.df_asin_keepa, allowMissingColumns=True)
        window = Window.partitionBy(['asin']).orderBy(F.asc("date_info"))
        self.df_save = self.df_save.withColumn(
            "rk", F.row_number().over(window=window)
        ).filter("rk=1").drop("rk").cache()
        self.df_save = self.df_save.select("asin", "date_info")
        self.df_save = self.df_save.withColumn("site_name", F.lit(self.site_name))
        self.df_save = self.df_save.withColumn("date_type", F.lit(self.date_type))
        if is_saturday_flag:
            self.df_save = self.df_save.withColumn("date_info", F.lit(self.date_info))
        else:
            self.df_save = self.df_save.filter(f"date_info='{self.date_info}'")

        self.df_save = self.df_save.join(
            df_result, on=['asin'], how='left'
        ).join(
            df_union_cate, on=['asin'], how='left'
        )
        self.df_save.show(10, truncate=False)
        # print(self.df_save.count())
        # self.df_save.groupBy("date_info").count().orderBy("count", ascending=False).show(truncate=False)
        # self.df_save.groupBy("asin_cate_flag").count().orderBy("count", ascending=False).show(truncate=False)
        # exit()
        print("self.df_save.count:", self.df_save.count())
        users = ["fangxingjun", "chenyuanjie", "pengyanbing"]

        title = f"日asin整合: {self.site_name}, {self.date_type}, {self.date_info}"
        content = f"整合asin完成--等待导出到pg提供爬虫使用--数量: {self.df_save.count()}"
        if self.date_type == 'month':
            title = f"月asin整合用于月流程流量选品的榜单计算: {self.site_name}, {self.date_type}, {self.date_info}"
            content = f"整合asin完成(爬虫无需抓取)--数量: {self.df_save.count()}"
        CommonUtil().send_wx_msg(users=users, title=title, content=content)

        if self.date_type in ('day', 'month'):
            # 写入 Doris dwd.dwd_asin_source_flag（LIST 三元组分区 site_name + date_type + date_info，全保留）
            # day  分区：date_info=yyyy-MM-dd
            # month 分区：date_info=yyyy-MM
            print(f"写入 Doris dwd.dwd_asin_source_flag, date_type={self.date_type}, date_info={self.date_info}")
            df_to_doris = self.df_save_asin_cate.select(
                F.col("site_name"),
                F.lit(self.date_type).alias('date_type'),
                F.col("date_info"),
                F.col("asin"),
                F.to_json(
                    F.expr("transform(split(asin_cate_flag, ','), x -> cast(x as int))")
                ).alias("asin_cate_flag"),
                F.col("bsr_latest_date").cast("date").alias("bsr_latest_date"),
                "bsr_30day_count",
                F.col("nsr_latest_date").cast("date").alias("nsr_latest_date"),
                "nsr_30day_count",
            )
            table_columns = "site_name, date_type, date_info, asin, asin_cate_flag, bsr_latest_date, bsr_30day_count, nsr_latest_date, nsr_30day_count"
            DorisHelper.spark_export_with_columns(
                df_save=df_to_doris,
                db_name='dwd',
                table_name='dwd_asin_source_flag',
                table_columns=table_columns,
            )
            print("Doris dwd.dwd_asin_source_flag 写入完毕")

    # def save_data(self):
    #     pass



if __name__ == '__main__':
    site_name = sys.argv[1]  # 参数1：站点
    date_type = sys.argv[2]  # 参数2：类型：week/4_week/month/quarter/day
    date_info = sys.argv[3]  # 参数3：年-周/年-月/年-季/年-月-日, 比如: 2022-1
    handle_obj = DwdNsrBsrKeepaAsin(site_name=site_name, date_type=date_type, date_info=date_info)
    handle_obj.run()
