import os
import sys

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

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


class DwtFlowKeepaAsin(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.df_flow_asin = self.spark.sql(f"select 1+1;")
        self.df_category_id = self.spark.sql(f"select 1+1;")
        self.df_keepa_asin = self.spark.sql(f"select 1+1;")
        self.df_calc_asin = self.spark.sql(f"select 1+1;")
        self.df_export_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):
        # 读取流量选品月asin
        sql = f"""
            select asin, asin_price as price, category_first_id, date_info as source_month
            from dwt_flow_asin
            where site_name = '{self.site_name}'
              and date_type = 'month'
              and date_info >= '2025-05'
              and asin_price is not null
              and asin_price > 0
        """
        df_flow_asin_month = self.spark.sql(sqlQuery=sql) \
            .withColumn('price', F.round(F.col('price'), 2).cast('decimal(10,2)'))

        # 读取ES最近30天缺少利润率的asin
        days_30_ago = (datetime.now() - timedelta(days=30)).strftime("%Y-%m-%d 00:00:00")
        es_read_options = {
            "es.nodes": EsUtils.__es_ip__,
            "es.port": EsUtils.__es_port__,
            "es.net.http.auth.user": EsUtils.__es_user__,
            "es.net.http.auth.pass": EsUtils.__es_passwd__,
            "es.nodes.wan.only": "false",
            "es.mapping.date.rich": "false",
            "es.scroll.size": "2000",
            "es.read.field.include": "asin,price,category_first_id,asin_crawl_date",
            "es.query": f'{{"query":{{"bool":{{"must":[{{"range":{{"price":{{"gt":0}}}}}},{{"range":{{"asin_crawl_date":{{"gte":"{days_30_ago}"}}}}}}],"must_not":{{"exists":{{"field":"profit_rate_extra.ocean_profit"}}}}}}}}}}'
        }
        df_flow_asin_30day = self.spark.read.format("org.elasticsearch.spark.sql") \
            .options(**es_read_options) \
            .load(f"{self.site_name}_flow_asin_30day") \
            .withColumn('price', F.round(F.col('price'), 2).cast('decimal(10,2)')) \
            .withColumn('source_month', F.date_format(F.col('asin_crawl_date'), 'yyyy-MM')) \
            .select('asin', 'price', 'category_first_id', 'source_month')

        # 合并两部分，按(asin, price)去重保留最新source_month
        self.df_flow_asin = df_flow_asin_month.union(df_flow_asin_30day).repartition(40, 'asin')
        window = Window.partitionBy(['asin', 'price']).orderBy(F.col('source_month').desc_nulls_last())
        self.df_flow_asin = self.df_flow_asin.withColumn(
            'rank', F.row_number().over(window=window)
        ).filter('rank = 1').drop('rank').cache()

        # 读取分类数据
        sql = f"""
        select category_first_id, en_name as category from dim_bsr_category_tree where site_name = '{self.site_name}' and nodes_num = 2
        """
        self.df_category_id = self.spark.sql(sqlQuery=sql).cache()

        # 读取keepa数据
        sql = f"""
        select asin, package_length, package_width, package_height, weight from dim_keepa_asin_info where site_name = '{self.site_name}' and date_info = '{self.date_info}'
        """
        self.df_keepa_asin = self.spark.sql(sqlQuery=sql).repartition(40, 'asin').filter(
            (F.col("package_length") > 0) & (F.col("package_width") > 0) & (F.col("package_height") > 0) & (F.col("weight") > 0)
        ).cache()

        # 读取已经计算过利润率的asin
        sql = f"""
        select asin, price, package_length, package_width, package_height, weight from dim_asin_profit_rate_info where site_name = '{self.site_name}' and date_info = '{self.date_info}'
        """
        self.df_calc_asin = self.spark.sql(sqlQuery=sql).repartition(40, 'asin').cache()

        # 读取已经导出过asin+price，避免重复计算
        sql = f"""
        select asin, price, package_length, package_width, package_height, weight from dwt_flow_keepa_asin where site_name = '{self.site_name}'
        """
        self.df_export_asin = self.spark.sql(sqlQuery=sql).repartition(40, 'asin').cache()

    def handle_data(self):
        self.df_save = self.df_flow_asin.join(
            self.df_category_id, on='category_first_id', how='left'
        ).join(
            self.df_keepa_asin, on='asin', how='inner'
        ).join(
            self.df_calc_asin, on=['asin', 'price', 'package_length', 'package_width', 'package_height', 'weight'], how='left_anti'
        ).join(
            self.df_export_asin, on=['asin', 'price', 'package_length', 'package_width', 'package_height', 'weight'], how='left_anti'
        ).cache()
        self.df_flow_asin.unpersist()
        self.df_category_id.unpersist()
        self.df_keepa_asin.unpersist()
        self.df_calc_asin.unpersist()
        self.df_export_asin.unpersist()

        start_key = 1
        self.df_save = self.df_save.withColumn(
            'part_key', F.ntile(50).over(Window.orderBy(F.rand())) + (start_key - 1)
        ).select(
            F.col('asin'),
            F.col('price'),
            F.col('category'),
            F.col('package_length'),
            F.col('package_width'),
            F.col('package_height'),
            F.col('weight'),
            F.col('part_key'),
            F.col('source_month'),
            F.lit(self.site_name).alias('site_name'),
            F.lit(self.date_info).alias('date_info')
        ).repartition(10)

    def save_data(self):
        hive_tb = "dwt_flow_keepa_asin"
        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_file_in_folder(hdfs_path)
        print(f"正在进行数据存储，当前存储的表名为：{hive_tb}，存储路径：{hdfs_path}")
        self.df_save.write.saveAsTable(name=hive_tb, format='hive', mode='append', partitionBy=partition_by)
        print("success")


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