import os
import sys

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

from utils.spark_util import SparkUtil
from pyspark.sql import functions as F, Window
from utils.common_util import CommonUtil
from utils.hdfs_utils import HdfsUtils
from datetime import datetime, timedelta


class DimAsinProfitRateInfo(object):

    def __init__(self, site_name, date_info):
        self.site_name = site_name
        self.date_info = date_info
        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}")

        self.df_asin_profit = self.spark.sql(f"select 1+1;")
        self.df_keepa_asin = self.spark.sql(f"select 1+1;")
        self.df_save = self.spark.sql(f"select 1+1;")

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

    def read_data(self):
        # 清洗整合利润率表
        sql = f"""
            select asin, price, category, ocean_profit, air_profit, package_length, package_width, package_height, weight, updated_time, date_info
            from dim_asin_profit_rate_info where site_name = '{self.site_name}';
        """
        self.df_asin_profit = self.spark.sql(sqlQuery=sql).repartition(40, 'asin').cache()

        # 读取keepa数据
        sql = f"""
            select asin, package_length, package_width, package_height, weight 
            from dim_keepa_asin_info where site_name = '{self.site_name}';
        """
        self.df_keepa_asin = self.spark.sql(sqlQuery=sql).repartition(40, 'asin')

    def handle_data(self):
        # 因为keepa数据存在更新的情况，保留与keepa最新数据所对应的数据行
        self.df_asin_profit = self.df_asin_profit.join(
            self.df_keepa_asin, on=['asin', 'package_length', 'package_width', 'package_height', 'weight'], how='inner'
        )
        # 去重
        window = Window.partitionBy(['asin', 'price']).orderBy(
            self.df_asin_profit.updated_time.desc_nulls_last()
        )
        self.df_asin_profit = self.df_asin_profit.withColumn(
            'rank', F.row_number().over(window=window)
        ).filter('rank = 1').drop('rank', 'date_info')

    def save_data(self):
        self.df_save = self.df_asin_profit.withColumn(
            "site_name", F.lit(self.site_name)
        ).withColumn(
            "date_info", F.lit(self.date_info)
        ).repartition(10).cache()

        new_count = self.df_save.count()
        print(f"最新数据量：{new_count}")

        hive_tb = "dim_asin_profit_rate_info"
        partition_dict = {
            "site_name": self.site_name,
            "date_info": self.date_info
        }
        partition_by = list(partition_dict.keys())
        hdfs_path = CommonUtil.build_hdfs_path(hive_tb, partition_dict)
        HdfsUtils.delete_hdfs_file(hdfs_path)
        print(f"正在进行数据存储，当前存储的表名为：{hive_tb}，存储路径：{hdfs_path}")
        self.df_save.write.saveAsTable(name=hive_tb, format='hive', mode='append', partitionBy=partition_by)

        # 验证实际写入数量，确保写入成功后再删除历史分区
        written_count = self.spark.sql(f"""
            select count(1) as cnt from {hive_tb}
            where site_name='{self.site_name}' and date_info='{self.date_info}'
        """).collect()[0]['cnt']
        print(f"实际写入数量：{written_count}，预期：{new_count}")
        if written_count != new_count:
            raise RuntimeError(
                f"写入数量校验失败！实际写入 {written_count} != 预期 {new_count}，终止删除历史分区，请人工检查。"
            )
        print("写入校验通过！")

        print(f"正在删除历史分区数据")
        self.spark.sql(f"""
        ALTER TABLE {hive_tb} DROP IF EXISTS PARTITION (site_name='{self.site_name}', date_info='{self.last_date_info}')
        """)
        HdfsUtils.delete_hdfs_file(
            CommonUtil.build_hdfs_path(hive_tb, partition_dict={"site_name": self.site_name, "date_info": self.last_date_info})
        )
        print("success!")


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