"""
数据源:
    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

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


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_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.date_type = self.judge_date_type()
        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.site_name == 'us' 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)

    def judge_date_type(self):
        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爬虫表
        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()
        df.show(10, truncate=False)
        print(f"sql: {sql}, content: {content}, count: {df.count()}")
        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_info == '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}'"
        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}"
        self.df_asin_from_day = self.read_data_common(sql=sql_asin_from_day, content="1.3. 读取dwd_nsr_bsr_keepa_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爬虫表
        year, month = self.date_info.split("-")
        sql_asin_syn = f"select asin from us_all_syn_st_month_{year} where date_info='{self.date_info}'"
        pdf_asin = self.engine_pg14.read_sql(sql_asin_syn)
        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)

    def handle_data(self):
        if self.date_type in ['month']:
            self.df_save = self.df_asin_from_st.unionByName(
                self.df_asin_from_adv, allowMissingColumns=True
            ).unionByName(
                self.df_asin_from_day, allowMissingColumns=True
            )
        if self.date_type == 'month_week':
            self.df_save = self.df_asin_from_st.unionByName(
                self.df_asin_from_adv, allowMissingColumns=True
            )
        if self.date_type == 'day':
            self.df_save = self.df_asin_from_day

        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'
        )
        # 处理同步逻辑
        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()}")

    # 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()
