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
from utils.common_util import CommonUtil
from datetime import datetime, timedelta

class EsAsinProfitRate(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__}")

        # ES相关配置
        self.es_client = EsUtils.get_es_client()
        self.es_profit_rate_index = f"{self.site_name}_profit_rate_extra_v2"
        self.es_profit_rate_body = self.get_es_profit_rate_body()
        self.es_profit_rate_options = self.get_es_profit_rate_options(self.es_profit_rate_index)

        self.df_asin_profit_rate = self.spark.sql(f"select 1+1;")

    @staticmethod
    def get_es_profit_rate_body():
        return {
            "settings": {
                "number_of_shards": "3",
                "number_of_replicas": "1"
            },
            "mappings": {
                "properties": {
                    "profit_key": {
                        "type": "keyword"
                    },
                    "asin": {
                        "type": "keyword"
                    },
                    "price": {
                        "type": "float"
                    },
                    "ocean_profit": {
                        "type": "float"
                    },
                    "air_profit": {
                        "type": "float"
                    },
                    "update_time": {
                        "type": "date",
                        "format": "yyyy-MM-dd"
                    }
                }
            }
        }

    @staticmethod
    def get_es_profit_rate_options(index_name):
        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": "profit_key",
            "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"
        }

    def run(self):
        self.save_profit_rate_add()
        self.update_history_index()

    def save_profit_rate_add(self):
        # 读取利润率整合数据（增量数据）
        sql = f"""
        select asin, price, ocean_profit, air_profit, updated_time from dim_asin_profit_rate_info
        where site_name = '{self.site_name}' and updated_time >= '{self.last_date_info}'
        """
        self.df_asin_profit_rate = self.spark.sql(sqlQuery=sql).repartition(40, 'asin')
        self.df_asin_profit_rate = self.df_asin_profit_rate.withColumn(
            'profit_key', F.concat_ws("_", F.col("asin"), F.col("price"))
        ).withColumn(
            "update_time",
            F.when(
                F.col("updated_time").isNotNull(),
                F.substring(F.col("updated_time"), 1, 10)
            ).otherwise(F.lit("1970-01-01"))
        ).select(
            'profit_key', 'asin', 'price', 'ocean_profit', 'air_profit', 'update_time'
        ).cache()
        print(f"增量利润率数据如下：")
        self.df_asin_profit_rate.show(10, False)

        print(f"创建利润率索引：{self.es_profit_rate_index}！")
        EsUtils.create_index(self.es_profit_rate_index, self.es_client, self.es_profit_rate_body)
        try:
            self.df_asin_profit_rate.write.format("org.elasticsearch.spark.sql") \
                .options(**self.es_profit_rate_options) \
                .mode("append") \
                .save()
            print(f"ES {self.es_profit_rate_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_profit_rate_index}')

    def update_history_index(self):
        """更新历史月度索引的利润率数据"""
        # 从 2025-05 开始，遍历到最新索引
        start_date = datetime(2025, 5, 1)
        current_date = start_date

        while True:
            year = current_date.year
            month = current_date.month
            month_str = f"{year}-{month:02d}"
            index_name = f"{self.site_name}_st_detail_month_{year}_{month:02d}"

            # 检查索引是否存在
            if not self.es_client.indices.exists(index=index_name):
                print(f"索引 {index_name} 不存在，停止遍历")
                break

            print(f"\n{'='*60}")
            print(f"开始处理索引: {index_name}")
            print(f"{'='*60}")

            try:
                self.update_single_history_index(index_name, month_str)
            except Exception as e:
                print(f"更新索引 {index_name} 失败: {str(e)}")

            # 移动到下个月
            if month == 12:
                current_date = datetime(year + 1, 1, 1)
            else:
                current_date = datetime(year, month + 1, 1)

    def update_single_history_index(self, index_name, month_str):
        """更新单个历史索引"""
        hive_sql = f"""
        select asin, asin_price as price from dwt_flow_asin where site_name = '{self.site_name}' and date_type = 'month' 
        and date_info = '{month_str}' and asin_price is not null
        """
        df_hive = self.spark.sql(hive_sql)

        df_update = self.df_asin_profit_rate.join(
            df_hive, on=['asin', 'price'], how='inner'
        ).withColumn(
            "profit_rate_extra",
            F.struct(
                F.col("ocean_profit").alias("ocean_profit"),
                F.col("air_profit").alias("air_profit")
            )
        ).select("asin", "profit_rate_extra")

        es_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.mapping.id": "asin",
            "es.resource": f"{index_name}/_doc",
            "es.batch.write.refresh": "false",
            "es.batch.size.entries": "5000",           # 批次数据量
            "es.write.operation": "update",
            "es.batch.write.concurrency": "5",         # 降低并发数，默认是自动（较高）
            "es.batch.write.retry.count": "3",         # 重试次数
            "es.batch.write.retry.wait": "30s",        # 重试等待
            "es.http.timeout": "5m",                   # 增加超时时间
            "es.internal.es.version.ignore": "true"    # 忽略版本检查
        }

        print(f"索引 {index_name} 待更新数据量: {df_update.count()}")
        df_update.show(5, False)
        df_update.repartition(10).write.format("org.elasticsearch.spark.sql") \
            .options(**es_options) \
            .mode("append") \
            .save()

        print(f"索引 {index_name} 更新完毕！")


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