import os
import sys

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

from utils.common_util import CommonUtil
from utils.spark_util import SparkUtil
from utils.es_util import EsUtils
from utils.db_util import DBUtil
from datetime import datetime
from pyspark.sql import functions as F


class EsAiAsinAdd(object):

    def __init__(self, site_name, date_type, date_info):
        self.site_name = site_name
        self.date_type = date_type
        self.date_info = date_info
        self.spark = SparkUtil.get_spark_session(f"{self.__class__.__name__}")

        self.pg_conn = DBUtil.get_connection_info("postgresql", "us")
        self.export_pg_tb = f"{self.site_name}_ai_asin_detail_month_{self.date_info.replace('-', '_')}"

        self.es_client = EsUtils.get_es_client()
        self.es_index = f"{self.site_name}_ai_asin_analyze_detail_{self.date_info.replace('-', '_')}"
        self.es_pipeline = f"{self.site_name}_ai_analyze_pipeline"
        self.es_options = self.get_es_options(self.es_index, self.es_pipeline)

        self.df_ai_asin_detail = self.spark.sql(f"select 1+1;")
        self.df_ai_asin_analyze = self.spark.sql(f"select 1+1;")
        self.df_profit_rate = self.spark.sql(f"select 1+1;")
        self.df_save_pg = self.spark.sql(f"select 1+1;")
        self.df_save_es = self.spark.sql(f"select 1+1;")

    @staticmethod
    def get_es_options(index_name, pipeline_id):
        return {
            "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.mapping.id": "asin",
            "es.resource": f"{index_name}/_doc",
            "es.batch.write.refresh": "false",
            "es.batch.write.retry.wait": "60s",
            "es.batch.size.entries": "5000",
            "es.nodes.wan.only": "false",
            "es.batch.write.concurrency": "40",
            "es.write.operation": "index",
            "es.ingest.pipeline": f"{pipeline_id}"
        }

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

    def read_data(self):
        # 读取asin信息库月数据
        sql1 = f"""
        select 
            site_name,
            asin,
            weight,
            bought_month,
            category,
            img,
            title,
            brand,
            account_name,
            account_addr,
            buy_box_seller_type,
            launch_time,
            img_num,
            variation_flag,
            variation_num,
            ao_val,
            category_id,
            category_current_id,
            parent_asin,
            bsr_rank,
            price,
            rating,
            total_comments,
            seller_id,
            fb_country_name,
            bought_month_mom,
            bought_month_yoy,
            is_new_flag,
            is_ascending_flag,
            review_json_list,
            launch_time_type,
            describe,
            product_json,
            product_detail_json
        from dwt_ai_asin_add 
        where site_name = '{self.site_name}' 
          and date_type = '{self.date_type}' 
          and date_info = '{self.date_info}'
        """
        self.df_ai_asin_detail = self.spark.sql(sqlQuery=sql1).repartition(40, 'asin').cache()
        print("ASIN信息库数据如下：")
        self.df_ai_asin_detail.show(10, True)

        # 读取Ai分析结果
        sql2 = f"""
        select
            asin, 					
            id as analyze_id,
            package_quantity,
            material,
            color,
            appearance,
            size,
            function,
            shape,
            scene_title,
            scene_comment,
            uses,
            theme,
            crowd,
            short_desc,
            title_pic_flag,
            title_word_flag,
            title_pic_content,
            title_word_content,
            array_to_string(package_quantity_arr, ',') as package_quantity_arr,
            package_quantity_flag,
            label_content,
            festival,
            multi_color_flag,
            multi_color_content
        from {self.site_name}_ai_asin_analyze_detail
        """
        self.df_ai_asin_analyze = SparkUtil.read_jdbc_query(
            session=self.spark,
            url=self.pg_conn["url"],
            pwd=self.pg_conn["pwd"],
            username=self.pg_conn["username"],
            query=sql2
        ).withColumn(
            'package_quantity_arr', F.split(F.col('package_quantity_arr'), ',')
        ).withColumn(
            'package_quantity_arr', F.expr('transform(package_quantity_arr, x -> cast(x as int))')
        ).repartition(40, 'asin').cache()
        print("AI分析数据如下：")
        self.df_ai_asin_analyze.show(10, True)

        # 读取利润率数据
        sql3 = f"""
        select asin, price, ocean_profit, air_profit
        from dim_asin_profit_rate_info
        where site_name = '{self.site_name}'
        """
        self.df_profit_rate = self.spark.sql(sql3).repartition(40, 'asin').cache()
        print("利润率数据如下：")
        self.df_profit_rate.show(10, True)

    def handle_data(self):
        self.df_save_pg = self.df_ai_asin_detail.join(
            self.df_ai_asin_analyze, 'asin', 'left_anti'
        ).select(
            'site_name', 'asin', 'weight', 'bought_month', 'category', 'img', 'title', 'brand', 'account_name', 
            'account_addr', 'buy_box_seller_type', 'launch_time', 'img_num', 'variation_flag', 'variation_num', 
            'ao_val', 'category_id', 'category_current_id', 'parent_asin', 'bsr_rank', 'price', 'rating', 
            'total_comments', 'seller_id', 'fb_country_name', 'review_json_list', 'launch_time_type', 'describe', 
            'product_json', 'product_detail_json', 'bought_month_mom', 'bought_month_yoy', 'is_new_flag', 
            'is_ascending_flag'
        )

        df_profit = self.df_profit_rate.withColumn(
            "profit_rate_extra",
            F.when(
                F.col("ocean_profit").isNull() & F.col("air_profit").isNull(),
                F.lit(None)
            ).otherwise(
                F.struct(
                    F.col("ocean_profit").alias("ocean_profit"),
                    F.col("air_profit").alias("air_profit")
                )
            )
        ).drop("ocean_profit", "air_profit")

        self.df_save_es = self.df_ai_asin_detail.join(
            self.df_ai_asin_analyze, 'asin', 'inner'
        ).withColumn(
            'profit_key', F.concat_ws("_", F.col("asin"), F.col("price"))
        ).join(
            df_profit, on=["asin", "price"], how="left"
        ).select(
            'account_addr', 'account_name', 'analyze_id', 'ao_val', 'appearance', 'asin', 'bought_month',
            'bought_month_mom', 'bought_month_yoy', 'brand', 'bsr_rank', 'buy_box_seller_type', 'category',
            'category_current_id', 'category_id', 'color', 'crowd', 'fb_country_name', 'festival', 'function',
            'img', 'img_num', 'is_ascending_flag', 'is_new_flag', 'label_content', 'launch_time', 'launch_time_type',
            'material', 'multi_color_content', 'multi_color_flag', 'package_quantity', 'package_quantity_arr',
            'package_quantity_flag', 'parent_asin', 'price', 'profit_key', 'profit_rate_extra', 'rating',
            'scene_comment', 'scene_title', 'seller_id', 'shape', 'short_desc', 'site_name', 'size', 'theme',
            'title', 'title_pic_content', 'title_pic_flag', 'title_word_content', 'title_word_flag',
            'total_comments', 'uses', 'variation_flag', 'variation_num', 'weight'
        )

    def create_pg_table(self):
        """
        根据模板表创建月度 PG 表：
          1. LIKE INCLUDING ALL：复制所有字段类型、其他列默认值、约束、索引
          2. 重建 id 列独立序列，避免与模板表共享同一序列
        """
        template_tb = "us_ai_asin_detail_month_2026_01"
        engine = DBUtil.get_db_engine("postgresql", "us")
        # 表已存在则跳过
        result = DBUtil.engine_exec_sql(engine, f"SELECT to_regclass('{self.export_pg_tb}')")
        if list(result)[0][0] is not None:
            print(f"PostgreSQL 表 {self.export_pg_tb} 已存在，跳过建表")
            return
        # 建表 + 为 id 列创建独立序列
        sql = f"""
            CREATE TABLE {self.export_pg_tb} (LIKE {template_tb} INCLUDING ALL);
            ALTER TABLE {self.export_pg_tb} ALTER COLUMN id DROP DEFAULT;
            CREATE SEQUENCE {self.export_pg_tb}_id_seq OWNED BY {self.export_pg_tb}.id;
            ALTER TABLE {self.export_pg_tb} ALTER COLUMN id SET DEFAULT nextval('{self.export_pg_tb}_id_seq')
        """
        DBUtil.exec_sql("postgresql", "us", sql)
        print(f"PostgreSQL 表 {self.export_pg_tb} 创建完成（独立自增序列）")

    def save_data(self):
        # 创建月度 PG 表
        self.create_pg_table()
        # 将新增asin导出给济苍
        try:
            self.df_save_pg.write.format("jdbc") \
                .option("url", self.pg_conn["url"]) \
                .option("dbtable", f"{self.export_pg_tb}") \
                .option("user", self.pg_conn["username"]) \
                .option("password", self.pg_conn["pwd"]) \
                .mode("append") \
                .save()
            CommonUtil.send_wx_msg(['wujicang', 'chenyuanjie'], 'ASIN信息库增量数据导出', f'详情：{self.export_pg_tb} {self.site_name} {self.date_type} {self.date_info}')
        except Exception as e:
            print("An error occurred while writing to PostgreSQL:", str(e))
            CommonUtil.send_wx_msg(['chenyuanjie'], '\u26A0 ASIN信息库增量数据导出失败', f'详情：{self.export_pg_tb} {self.site_name} {self.date_type} {self.date_info}')

        # 将增量asin导出到es
        try:
            EsUtils.create_index(self.es_index, self.es_client, EsUtils.get_es_ai_body())
            self.df_save_es.write.format("org.elasticsearch.spark.sql") \
                .options(**self.es_options) \
                .mode("append") \
                .save()
            print(f"ES {self.es_index} 索引更新完毕！")
        except Exception as e:
            print("An error occurred while writing to Elasticsearch:", str(e))
            CommonUtil.send_wx_msg(['chenyuanjie'], '\u26A0 ES数据更新失败', f'失败索引：{self.es_index}')


if __name__ == "__main__":
    site_name = sys.argv[1]
    date_type = sys.argv[2]
    date_info = sys.argv[3]
    print("开始执行时间：", datetime.now().strftime("%Y-%m-%d %H:%M"))
    handle_obj = EsAiAsinAdd(site_name, date_type, date_info)
    handle_obj.run()
    print("执行结束时间：", datetime.now().strftime("%Y-%m-%d %H:%M"))
    print("success！！！")
