"""
author: CT
description: 导出待计算利润率的 ASIN
             1) Hive dwt_flow_asin 月维度读取 date_info >= '2025-05' 的所有 ASIN：
                asin / price / category_first_id / asin_crawl_date
             2) Doris dwt.{site}_flow_asin_30day 读取所有相关 ASIN：
                asin / price / category_first_id / asin_crawl_date
             3) union 后按 (asin, price) 去重保留 asin_crawl_date 最新
             4) LEFT JOIN 分类、INNER JOIN keepa 增量
                keepa 表已整合为单分区快照，按 updated_time > last_date_info 筛增量
             5) keepa 关联到的 ASIN 全部导出 PG {site}_asin_profit_rate_calc 重新计算利润率
执行示例: spark-submit export_need_profit_rate.py us 2026-05-15
"""
import os
import sys
from datetime import datetime, timedelta

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

from pyspark.sql import functions as F, Window

from utils.spark_util import SparkUtil
from utils.db_util import DBUtil
from utils.DorisHelper import DorisHelper

START_MONTH = '2025-05'


class ExportNeedProfitRate(object):

    def __init__(self, site_name, date_info):
        self.site_name = site_name
        self.date_info = date_info  # 计算时间 yyyy-MM-dd
        # keepa 增量过滤下限：date_info - 1 天（如 date_info=2026-05-13 → last_date_info=2026-05-12）
        self.last_date_info = (datetime.strptime(date_info, "%Y-%m-%d").date() - timedelta(days=1)).strftime("%Y-%m-%d")
        self.spark = SparkUtil.get_spark_session(
            f"{self.__class__.__name__}: {self.site_name} {self.date_info}"
        )

    def run(self):
        df_export = self.build_export_df()
        self.write_to_pg(df_export)

    def build_export_df(self):
        # 1. Hive dwt_flow_asin 月维度，date_info >= 2025-05 所有月份
        sql_dwt = f"""
            SELECT asin,
                   asin_price       AS price,
                   category_first_id,
                   asin_crawl_date
            FROM dwt_flow_asin
            WHERE site_name = '{self.site_name}'
              AND date_type = 'month'
              AND date_info >= '{START_MONTH}'
              AND asin_price > 0
        """
        print(f"sql_dwt =\n{sql_dwt}")
        df_dwt = self.spark.sql(sqlQuery=sql_dwt) \
            .withColumn('price', F.round(F.col('price'), 2).cast('decimal(20,2)')) \
            .withColumn('asin_crawl_date', F.to_timestamp(F.col('asin_crawl_date')))

        # 2. Doris dwt.{site}_flow_asin_30day 所有相关 ASIN
        #    spark_import_with_connector 不支持 WHERE 下推，price > 0 过滤改到 Spark 端
        table_identifier = f"dwt.{self.site_name}_flow_asin_30day"
        read_fields = "asin,price,category_first_id,asin_crawl_date"
        df_doris = DorisHelper.spark_import_with_connector(self.spark, table_identifier, read_fields) \
            .filter(F.col('price') > 0) \
            .withColumn('price', F.round(F.col('price'), 2).cast('decimal(20,2)')) \
            .withColumn('asin_crawl_date', F.col('asin_crawl_date').cast('timestamp')) \
            .select('asin', 'price', 'category_first_id', 'asin_crawl_date')

        # 3. union + 按 (asin, price) 去重保留 asin_crawl_date 最新
        df_flow = df_dwt.unionByName(df_doris).repartition(40, 'asin', 'price')
        window = Window.partitionBy('asin', 'price').orderBy(F.col('asin_crawl_date').desc_nulls_last())
        df_flow = df_flow.withColumn('rk', F.row_number().over(window)) \
                         .filter('rk = 1') \
                         .drop('rk') \
                         .cache()

        # 4. 分类名 LEFT JOIN
        sql_cate = f"""
            SELECT category_first_id, en_name AS category
            FROM dim_bsr_category_tree
            WHERE site_name = '{self.site_name}' AND nodes_num = 2
        """
        df_cate = self.spark.sql(sqlQuery=sql_cate)

        # 5. keepa 当日增量 INNER JOIN
        # keepa 表已整合为单分区快照，用 updated_time > last_date_info 筛"近一天更新"的增量
        sql_keepa = f"""
            SELECT asin, package_length, package_width, package_height, item_weight AS weight
            FROM dim_keepa_asin_info
            WHERE site_name = '{self.site_name}' AND updated_time >= '{self.last_date_info}'
        """
        df_keepa = self.spark.sql(sqlQuery=sql_keepa) \
            .filter((F.col('package_length') > 0) &
                    (F.col('package_width') > 0) &
                    (F.col('package_height') > 0)) \
            .repartition(40, 'asin')

        df_result = df_flow \
            .join(df_cate, on='category_first_id', how='left') \
            .join(df_keepa, on='asin', how='inner') \
            .withColumn('source_month', F.date_format(F.col('asin_crawl_date'), 'yyyy-MM')) \
            .withColumn('part_key', F.ntile(50).over(Window.orderBy(F.rand()))) \
            .select(
                'asin', 'price', 'category',
                'package_length', 'package_width', 'package_height', 'weight',
                'part_key', 'source_month', 'asin_crawl_date',
            ).cache()

        count = df_result.count()
        print(f"待计算利润率数据量：{count:,}")
        df_result.show(10, truncate=False)
        df_flow.unpersist()
        return df_result

    def write_to_pg(self, df_export):
        con_info   = DBUtil.get_connection_info(db_type='postgresql_cluster', site_name=self.site_name)
        table_name = f"{self.site_name}_asin_profit_rate_calc"
        print(f"导出到 PG {table_name}")
        df_export.write.format("jdbc") \
            .option("url",      con_info["url"]) \
            .option("dbtable",  table_name) \
            .option("user",     con_info["username"]) \
            .option("password", con_info["pwd"]) \
            .mode("append") \
            .save()
        print("success")


if __name__ == "__main__":
    site_name = sys.argv[1]
    date_info = sys.argv[2]
    ExportNeedProfitRate(site_name, date_info).run()
