"""
author: CT
description: 利润率数据增量整合 — ODS → Hive → Doris
             前置：ods_asin_profit_rate.py 已将当日 PG 增量 sqoop 到 ods_asin_profit_rate
             步骤：
               1) 读 ods_asin_profit_rate 当日分区（原始增量）+ dim_asin_profit_rate_info 历史分区
               2) 按 (asin, price) 去重，updated_time desc 保留最新一行
               3) 覆盖 dim_asin_profit_rate_info 当日分区为整合后全量快照；校验后删除历史分区
               4) 当日 ODS 增量写入 Doris dwd.dwd_asin_profit_rate_latest
                  Doris UNIQUE KEY(site_name, asin, price) + sequence_col=update_time 自动取最新
执行示例: spark-submit dim_asin_profit_rate_info.py us 2026-05-15
"""
import os
import sys

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.common_util import CommonUtil
from utils.hdfs_utils import HdfsUtils
from utils.DorisHelper import DorisHelper


class DimAsinProfitRateInfo(object):

    def __init__(self, site_name, date_info):
        self.site_name   = site_name
        self.date_info   = date_info
        self.spark = SparkUtil.get_spark_session(
            f"{self.__class__.__name__}: {self.site_name} {self.date_info}"
        )
        self.ods_table   = "ods_asin_profit_rate"
        self.hive_table  = "dim_asin_profit_rate_info"
        self.doris_db    = "dwd"
        self.doris_table = "dwd_asin_profit_rate_latest"
        self.df_today    = None
        self.df_history  = None
        self.df_save     = None

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

    def read_data(self):
        """读当日 ODS 增量（ods_asin_profit_rate）+ dim 历史所有分区"""
        sql_today = f"""
            SELECT asin, price, category, ocean_profit, air_profit,
                   package_length, package_width, package_height, weight,
                   asin_crawl_date, updated_time
            FROM {self.ods_table}
            WHERE site_name = '{self.site_name}' AND date_info = '{self.date_info}'
        """
        print(f"sql_today =\n{sql_today}")
        # cache：save_data 会 DROP PARTITION 当日分区，write_to_doris 还要复用此 df，必须先物化
        # count() 强制立即物化，确保 cache 在 DROP/saveAsTable 之前已填充到 BlockManager
        self.df_today = self.spark.sql(sqlQuery=sql_today).repartition(40, 'asin', 'price').cache()
        self.df_today.count()

        sql_history = f"""
            SELECT asin, price, category, ocean_profit, air_profit,
                   package_length, package_width, package_height, weight,
                   asin_crawl_date, updated_time
            FROM {self.hive_table}
            WHERE site_name = '{self.site_name}' AND date_info != '{self.date_info}'
        """
        print(f"sql_history =\n{sql_history}")
        self.df_history = self.spark.sql(sqlQuery=sql_history).repartition(40, 'asin', 'price').cache()

    def handle_data(self):
        """union 当日 + 历史，按 (asin, price) 去重保留 updated_time desc 最新一行"""
        df_all = self.df_today.unionByName(self.df_history)
        window = Window.partitionBy('asin', 'price').orderBy(
            F.col('updated_time').desc_nulls_last(),
        )
        self.df_save = df_all.withColumn('rk', F.row_number().over(window)) \
                             .filter('rk = 1') \
                             .drop('rk')

    def save_data(self):
        """覆盖当日分区为整合后全量快照；写入校验通过后删除历史分区"""
        self.df_save = self.df_save \
            .withColumn('site_name', F.lit(self.site_name)) \
            .withColumn('date_info', F.lit(self.date_info)) \
            .repartition(40).cache()
        new_count = self.df_save.count()
        old_count = self.df_history.count()
        print(f"整合后 new_count={new_count:,}，历史 old_count={old_count:,}")

        partition_dict = {"site_name": self.site_name, "date_info": self.date_info}
        partition_by   = list(partition_dict.keys())
        hdfs_path      = CommonUtil.build_hdfs_path(self.hive_table, partition_dict)

        # 先清掉 sqoop 写入的当日分区（hdfs + metastore），避免重复 append
        print(f"清空当日分区 hdfs：{hdfs_path}")
        HdfsUtils.delete_hdfs_file(hdfs_path)
        self.spark.sql(
            f"ALTER TABLE {self.hive_table} DROP IF EXISTS "
            f"PARTITION (site_name='{self.site_name}', date_info='{self.date_info}')"
        )
        # 写入整合后的当日分区（try-except 保护：写入失败立即抛出，绝不进入删除历史分区流程）
        print(f"写入 Hive 当日分区：{hdfs_path}")
        try:
            self.df_save.write.saveAsTable(
                name=self.hive_table, format='hive', mode='append', partitionBy=partition_by
            )
        except Exception as e:
            print(f"❌ Hive 当日分区写入失败：{e}；保留所有历史分区不删，请人工排查！")
            self.df_history.unpersist()
            raise

        # 验证当日分区实际写入行数与预期一致（防止 saveAsTable 静默丢数据）
        written_count = self.spark.sql(
            f"SELECT COUNT(1) AS cnt FROM {self.hive_table} "
            f"WHERE site_name='{self.site_name}' AND date_info='{self.date_info}'"
        ).collect()[0]['cnt']
        print(f"实际写入行数 written_count={written_count:,}，预期 new_count={new_count:,}")
        if written_count != new_count:
            print(f"❌ 写入校验失败 written_count={written_count} != new_count={new_count}，保留历史分区不删！")
            self.df_history.unpersist()
            return

        # 安全控制：当日整合分区行数必须 >= 历史总和才允许删历史（防止整合丢失数据）
        if new_count < old_count:
            print(f"⚠️ 校验未通过 new_count={new_count} < old_count={old_count}，跳过历史分区删除，请人工排查！")
            self.df_history.unpersist()
            return

        print(f"校验通过，删除所有 date_info < {self.date_info} 的历史分区")
        partitions = self.spark.sql(f"SHOW PARTITIONS {self.hive_table}").collect()
        for row in partitions:
            # row[0] 形如 "site_name=us/date_info=2026-05-10"
            parts = dict(p.split('=') for p in row[0].split('/'))
            p_site = parts.get('site_name')
            p_date = parts.get('date_info')
            if p_site == self.site_name and p_date < self.date_info:
                print(f"  删除分区: site_name={p_site}, date_info={p_date}")
                self.spark.sql(
                    f"ALTER TABLE {self.hive_table} DROP IF EXISTS "
                    f"PARTITION (site_name='{p_site}', date_info='{p_date}')"
                )
                HdfsUtils.delete_hdfs_file(
                    CommonUtil.build_hdfs_path(
                        self.hive_table,
                        partition_dict={"site_name": p_site, "date_info": p_date},
                    )
                )
        print("历史分区清理完毕!")
        self.df_history.unpersist()

    def write_to_doris(self):
        """当日 sqoop 增量数据写 Doris dwd_asin_profit_rate_latest（不回灌历史）
        Doris UNIQUE KEY(site_name, asin, price) + sequence_col=update_time 自动按更新时间取最新
        """
        df_to_doris = self.df_today.select(
            F.lit(self.site_name).alias('site_name'),
            F.col('asin'),
            F.round(F.col('price'), 2).cast('decimal(20,2)').alias('price'),
            F.col('ocean_profit'),
            F.col('air_profit'),
            F.to_timestamp(F.col('asin_crawl_date')).alias('asin_crawl_date'),
            F.to_timestamp(F.col('updated_time')).alias('update_time'),
        ).cache()
        count = df_to_doris.count()
        print(f"写入 Doris 增量数据量：{count:,}")
        df_to_doris.show(10, truncate=False)

        table_columns = "site_name, asin, price, ocean_profit, air_profit, asin_crawl_date, update_time"
        DorisHelper.spark_export_with_columns(
            df_save=df_to_doris,
            db_name=self.doris_db,
            table_name=self.doris_table,
            table_columns=table_columns,
        )
        df_to_doris.unpersist()
        self.df_today.unpersist()
        self.df_save.unpersist()
        print("success!")


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