"""
数据源:
    1. 搜索词对应的市场asin
    2. asin详情页对应的广告位关联asin
    3. 当前月对应所有天的新品和榜单和keep的asin
流程:
    month: 新增未抓取的asin, 123
    month_week: 清空爬虫的syn表, 12
    day: 3
"""

import os
import random
import sys
import time
import traceback
from datetime import datetime

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 utils.secure_db_client import get_remote_engine
from utils.spark_util import SparkUtil
from datetime import datetime
from dateutil.relativedelta import relativedelta


class DwtAsinSync(Templates):

    def __init__(self, site_name="us", date_type="month", date_info="2026-01"):
        super().__init__()
        self.site_name = site_name
        self.date_type = date_type
        self.date_info = date_info
        self.db_save = f"dwt_asin_sync"
        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_asin_flow = self.spark.sql(f"select 1+1;")
        self.df_asin_from_st = self.spark.sql(f"select 1+1;")
        self.df_asin_from_adv = self.spark.sql(f"select 1+1;")
        self.df_asin_from_day = self.spark.sql(f"select 1+1;")
        self.df_asin_variation = self.spark.sql(f"select 1+1;")
        self.df_asin_stable = self.spark.sql(f"select 1+1;")
        self.df_asin_syn = self.spark.sql(f"select 1+1;")
        self.engine_pg14 = get_remote_engine(
            site_name=self.site_name,
            db_type='postgresql_14'
        )
        self.engine_mysql = get_remote_engine(
            site_name=self.site_name,
            db_type='mysql'
        )
        print(type(self.engine_mysql))
        print(self.engine_mysql)
        self.get_date_info_tuple()  # self.date_info_tuple
        self.table_syn = f"{self.site_name}_all_syn_st_month_{self.date_info.replace('-', '_')}" if self.date_type == 'month' else f"{self.site_name}_all_syn_st_{self.date_info.replace('-', '_')}"
        self.partitions_by = ['site_name', 'date_type', 'date_info']
        self.reset_partitions(partitions_num=5)
        self.date_type = self.judge_date_type()
        # 导出到利润表
        self.pg_table = f"{site_name}_asin_profit_keepa_add"

    @staticmethod
    def judge_today():
        today = datetime.now().day
        if 10 <= today <= 25:
            return True
        else:
            return False

    @staticmethod
    def get_prev_month(month_str: str) -> str:
        """
        输入: YYYY-MM
        输出: 上一个月，格式仍为 YYYY-MM

        示例:
        get_prev_month("2026-01") -> "2025-12"
        """
        dt = datetime.strptime(month_str, "%Y-%m")
        return (dt - relativedelta(months=1)).strftime("%Y-%m")

    def judge_date_type(self):
        day_flag = self.judge_today()
        print(f"site_name: {self.site_name}, date_type: {self.date_type}, date_info: {self.date_info}")
        if self.date_type in ["month", "month_week"]:
            year, month = self.date_info.split("-")
            sql = f"select count(*) as st_count from {self.site_name}_brand_analytics_month_{year} where year={year} and month={int(month)} ;"
            df = self.engine_mysql.read_sql(sql)
            print("sql:", sql, df.shape)
            if list(df.st_count)[0] >= 100_0000:
                self.date_type = "month"  # 追加syn爬虫表
            else:
                # self.date_type = "month_week"  # 清空syn爬虫表
                while True:
                    try:
                        if day_flag:
                            # 1. 只有属于10-25号期间执行才清空syn爬虫表; 2. 月搜索词没有导入进来
                            sql_truncate = f"truncate {self.table_syn};"
                            print(f"月搜索词没有导入进来, 需要先清空表, sql_truncate: {sql_truncate}")
                            self.engine_pg14.execute(sql_truncate)
                        break
                    except Exception as e:
                        print(e, traceback.format_exc())
                        time.sleep(random.randint(3, 10))
                        self.engine_pg14 = get_remote_engine(
                            site_name=self.site_name,
                            db_type='postgresql_14'
                        )
                        continue
        return self.date_type

    def read_data_common(self, sql, content, col_dup='asin'):
        df = self.spark.sql(sqlQuery=sql).drop_duplicates([col_dup]).cache()
        print(f"sql: {sql}, content: {content}, count: {df.count()}")
        df.show(10, truncate=False)
        return df

    def read_data(self):
        # 读取asin数据源
        sql_asin_from_st = f"select asin, 1 as asin_flag from dim_st_asin_info where site_name='{self.site_name}' and date_type='{self.date_type}' and date_info='{self.date_info}'"
        self.df_asin_from_st = self.read_data_common(sql=sql_asin_from_st, content="1.1. 读取dim_st_asin_info表的st-asin数据")

        sql_asin_from_adv = f"select related_asin from dwt_asin_related_traffic where site_name='{self.site_name}' and date_type='month_week' and date_info='{self.date_info}'"
        self.df_asin_from_adv = self.read_data_common(sql=sql_asin_from_adv, content="1.2. 读取dwt_asin_related_traffic表的asin广告位数据", col_dup="related_asin")
        self.df_asin_from_adv = self.df_asin_from_adv.withColumn(
            'asin', F.explode(F.split(F.col('related_asin'), ','))
        ).select('asin').drop_duplicates(['asin'])
        self.df_asin_from_adv = self.df_asin_from_adv.withColumn('asin_flag', F.lit(2))

        if self.date_type == 'day':
            sql_asin_from_day = f"select asin, 3 as asin_flag from dwd_nsr_bsr_keepa_asin where site_name='{self.site_name}' and date_type='{self.date_type}' and date_info='{self.date_info}' and asin_type!='1'"
        else:
            sql_asin_from_day = f"select asin, 3 as asin_flag from dwd_nsr_bsr_keepa_asin where site_name='{self.site_name}' and date_type='day' and date_info in {self.date_info_tuple} and asin_type!='1'"
        self.df_asin_from_day = self.read_data_common(sql=sql_asin_from_day, content="1.3. 读取dwd_nsr_bsr_keepa_asin表的asin新品和榜单数据")
        # 获取上一个月的流量选品月份数据, amazon月销>0的asin
        prev_month = self.get_prev_month(self.date_info)
        sql_asin_flow = f"select asin, 4 as asin_flag from dwt_flow_asin where site_name='{self.site_name}' and date_type='{self.date_type}' and date_info='{prev_month}' and asin_bought_month>0"
        self.df_asin_flow = self.read_data_common(sql=sql_asin_flow, content="1.4. 读取dwt_flow_asin表的asin数据")

        # 读取asin属性值
        sql_asin_variation = f"""select asin, 1 as asin_is_variation from dim_asin_variation_info where site_name="{self.site_name}";"""
        self.df_asin_variation = self.read_data_common(sql=sql_asin_variation, content="2.1 读取dim_asin_variation_info表的asin变体属性")
        sql_asin_stable = f"""select asin, asin_volume as volume, asin_weight_str as weight_str from dim_asin_stable_info where site_name="{self.site_name}";"""
        self.df_asin_stable = self.read_data_common(sql=sql_asin_stable, content="2.2 读取dim_asin_variation_info表的asin重量体积属性")
        # 读取syn爬虫表
        while True:
            try:
                table_syn = f"{self.site_name}_all_syn_st_day_{self.date_info.replace('-', '_')}" if self.date_type == 'day' else f"{self.site_name}_all_syn_st_month_{self.date_info.replace('-', '_')}"
                sql_asin_syn = f"select asin from {table_syn};"
                pdf_asin = self.engine_pg14.read_sql(sql_asin_syn)
                print((f"pdf_asin: {pdf_asin.shape}, sql_asin_syn: {sql_asin_syn}"))
                if pdf_asin.shape[0]:
                    schema = StructType([
                        StructField('asin', StringType(), True),
                    ])
                    self.df_asin_syn = self.spark.createDataFrame(pdf_asin, schema=schema)
                    self.df_asin_syn = self.df_asin_syn.drop_duplicates(["asin"]).cache()
                    print(f"self.df_asin_syn: {self.df_asin_syn.count()}")
                    self.df_asin_syn.show(10, truncate=False)
                break
            except Exception as e:
                time.sleep(100)
                self.engine_pg14 = get_remote_engine(
                    site_name=self.site_name,
                    db_type='postgresql_14'
                )
                self.engine_mysql = get_remote_engine(
                    site_name=self.site_name,
                    db_type='mysql'
                )
                continue

    def sava_profit_data(self):
        df = self.df_asin_from_day.select("asin")
        df = df.drop_duplicates(['asin'])
        df = df.withColumn("month", F.lit(self.date_info[:7]))
        # 排除 dim_keepa_asin_info 中已有有效keepa数据的ASIN
        # 若 package_length/width/height/weight 任意一个 < 0，视为数据异常，不排除（需重新抓取）
        print("8. 排除已有keepa数据的ASIN (dim_keepa_asin_info)")
        df_keepa = self.spark.sql(f"""
                    select asin from dim_keepa_asin_info
                    where site_name      = '{self.site_name}'
                      and package_length >= 0
                      and package_width  >= 0
                      and package_height >= 0
                      and weight         >= 0
                """).repartition(40, 'asin')
        df = df.join(df_keepa, on='asin', how='left_anti').cache()
        print(f"排除keepa后数据量: {df.count()}")

        # 排除 {pg_table} 中已导出的ASIN
        print(f"9. 排除已导出的ASIN ({self.pg_table})")
        pg_con_info = DBUtil.get_connection_info("postgresql_cluster", self.site_name)
        df_exported = SparkUtil.read_jdbc_query(
            session=self.spark,
            url=pg_con_info['url'],
            username=pg_con_info['username'],
            pwd=pg_con_info['pwd'],
            query=f"select asin from {self.pg_table}"
        ).repartition(40, 'asin')
        df = df.join(df_exported, on='asin', how='left_anti').cache()
        print(f"排除已导出后数据量: {df.count()}")

        df.show(10, truncate=False)
        total = df.count()
        print(f"10. 写入 PostgreSQL 表 {self.pg_table}，共 {total} 条")
        # quit()
        df = df.select(
            F.col("asin"),
            F.col("month"),
        ).cache()
        con_info = DBUtil.get_connection_info('postgresql_cluster', self.site_name)
        df.write.format("jdbc") \
            .option("url", con_info["url"]) \
            .option("dbtable", self.pg_table) \
            .option("user", con_info["username"]) \
            .option("password", con_info["pwd"]) \
            .mode("append") \
            .save()
        print(f"写入完毕")

    def handle_data(self):
        if self.date_type in ['month']:
            self.df_save = self.df_asin_flow.unionByName(
                self.df_asin_from_st, allowMissingColumns=True
            ).unionByName(
                self.df_asin_from_adv, allowMissingColumns=True
            ).unionByName(
                self.df_asin_from_day, allowMissingColumns=True
            )
            # self.df_save = self.df_asin_from_adv.unionByName(
            #     self.df_asin_from_day, allowMissingColumns=True
            # )
        if self.date_type == 'month_week':
            self.df_save = self.df_asin_flow.unionByName(
                self.df_asin_from_st, allowMissingColumns=True
            ).unionByName(
                self.df_asin_from_adv, allowMissingColumns=True
            )
        if self.date_type == 'day':
            self.df_save = self.df_asin_from_day
            # 同步到利润表
            self.sava_profit_data()

        print(f"去重之前的asin_flag类型数量统计")
        self.df_save.groupby(f"asin_flag").agg(
            F.count("asin").alias("asin_flag_count")
        ).show()
        window = Window.partitionBy(['asin']).orderBy(F.desc(f"asin_flag"))
        self.df_save = self.df_save.withColumn(f"asin_flag_row_number", F.row_number().over(window)) \
            .filter(f'asin_flag_row_number = 1').drop(f"asin_flag_row_number")
        print(f"去重之后的asin_flag类型数量统计")
        self.df_save.groupby(f"asin_flag").agg(
            F.count("asin").alias("asin_flag_count")
        ).show()

        self.df_save = self.df_save.join(
            self.df_asin_variation, on='asin', how='left'
        ).join(
            self.df_asin_stable, on='asin', how='left'
        )
        # 处理同步逻辑
        print("==="*20)
        print(f"{type(self.df_asin_syn)}: {self.df_asin_syn.count()}")
        if self.date_type != 'day' and self.df_asin_syn.count() > 1:
            self.df_save = self.df_save.join(self.df_asin_syn, on=['asin'], how="left_anti")
        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))
        self.df_save = self.df_save.withColumn('date_info', F.lit(self.date_info))
        self.df_save.show(10, truncate=False)
        print(f"self.df_save: {self.df_save.count()}")
        users = ["fangxingjun", "chenyuanjie", "pengyanbing"]
        title = f"dwd_asin_to_pg: {self.site_name}, {self.date_type}, {self.date_info}"
        content = f"整合asin完成--等待导出到pg提供爬虫使用--self.df_save.count: {self.df_save.count()}"
        CommonUtil().send_wx_msg(users=users, title=title, content=content)

    # 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
    date_info = sys.argv[3]  # 参数3：年-周/年-月/年-季, 比如: 2022-1
    handle_obj = DwtAsinSync(site_name=site_name, date_type=date_type, date_info=date_info)
    handle_obj.run()
