import os
import sys

# =============================================================================
# 【ES写入异常排查笔记】EsHadoopNoNodesLeftException: all nodes failed
# =============================================================================
# 现象：Spark executor 写入ES时报 Connection error - all nodes failed，
#       tried [[210:9200, 217:9200, 216:9200]]
#
# 排查步骤：
#   在各 hadoop 节点执行：curl -u user:pwd http://{es_node}:9200
#   - 返回 401  → 网络可达（正常，未带认证）
#   - 连接超时  → 网络不可达
#
# 原因1：网络不可达（Executor 只能访问部分ES节点）
#   根因：ES扩容后新节点未对Spark集群开放网络权限，或节点间VLAN隔离
#   解法：设置 es.nodes.wan.only=true，强制所有请求只走配置节点，不做集群节点发现
#
# 原因2：ES节点过载（本项目实际原因）
#   根因：大数据量高并发写入，ES处理不过来拒绝连接，重试所有节点均失败
#   解法：es.nodes.wan.only=true 将所有请求集中到单个协调节点(217)，
#         由217内部分发到shard，相当于间接限流，降低对各数据节点的直接冲击
#         如仍报错，进一步降低并发和批次大小：
#           es.batch.size.entries: 2000
#           es.batch.write.concurrency: 5
#           es.batch.write.retry.count: 5
#           es.batch.write.retry.wait: 60s
#
# 结论：单入口访问ES集群时，es.nodes.wan.only=true 是推荐配置
# =============================================================================

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

from utils.spark_util import SparkUtil
from utils.es_util import EsUtils
from utils.DorisHelper import DorisHelper
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"
                    },
                    "asin_crawl_date": {
                        "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.read_profit_rate_add()
        # 利润率主索引
        self.save_profit_rate_to_es()
        # st_detail_month 近3个月
        for index_name in self.get_recent_indexes("st_detail_month"):
            date_info = self.get_date_info_from_index(index_name)
            df_asin = self.spark.sql(f"""
                select asin from dwt_flow_asin
                where site_name='{self.site_name}' and date_type='month' and date_info='{date_info}'
            """).repartition(40, 'asin')
            self.update_by_hive_asin(index_name, df_asin)
        # 年度信息库
        df_year_asin = self.spark.sql(f"""
            select asin from dwt_ai_asin_year
            where site_name='{self.site_name}'
              and date_info = (select max(date_info) from dwt_ai_asin_year where site_name='{self.site_name}')
        """).repartition(40, 'asin')
        self.update_by_hive_asin(f"{self.site_name}_ai_asin_analyze_detail_last365_day", df_year_asin)
        # ai_asin_analyze_detail 近3个月（从 ES 读取 asin+price，按 asin+price 关联）
        for index_name in self.get_recent_indexes("ai_asin_analyze_detail"):
            self.update_by_es_asin_price(index_name)
        # 30天流量选品（从 ES 读取 asin+price，按 asin+price 关联）
        self.update_by_es_asin_price(f"{self.site_name}_flow_asin_30day")

    def read_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(
            'price', F.round(F.col('price'), 2)
        ).withColumn(
            'profit_key', F.concat_ws("_", F.col("asin"), F.col("price"))
        ).withColumn(
            "update_time", F.substring(F.col("updated_time"), 1, 10)
        ).select(
            'profit_key', 'asin', 'price', 'ocean_profit', 'air_profit', 'update_time'
        )

        # 从Doris获取asin_crawl_date
        df_crawl_date = DorisHelper.spark_import_with_connector(
            session=self.spark,
            table_identifier=f"selection.{self.site_name}_asin_latest_detail",
            read_fields="asin,asin_crawl_date"
        ).withColumn(
            "asin_crawl_date", F.substring(F.col("asin_crawl_date"), 1, 10)
        ).repartition(40, 'asin')

        self.df_asin_profit_rate = self.df_asin_profit_rate.join(
            df_crawl_date, on='asin', how='left'
        ).select(
            'profit_key', 'asin', 'price', 'ocean_profit', 'air_profit', 'update_time', 'asin_crawl_date'
        ).cache()
        print(f"增量利润率数据如下：")
        self.df_asin_profit_rate.show(10, False)

    def save_profit_rate_to_es(self):
        print(f"\n{'='*60}")
        print(f"开始更新利润率索引：{self.es_profit_rate_index}")
        print(f"{'='*60}")
        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 get_recent_indexes(self, prefix, max_count=3):
        """从date_info对应月份开始往前找，返回最近max_count个存在的索引名"""
        indexes = []
        current = datetime.strptime(self.date_info, "%Y-%m-%d").replace(day=1)
        checked = 0
        while len(indexes) < max_count and checked < 24:
            index_name = f"{self.site_name}_{prefix}_{current.year}_{current.month:02d}"
            if EsUtils.exist_index(index_name, self.es_client):
                indexes.append(index_name)
            if current.month == 1:
                current = current.replace(year=current.year - 1, month=12)
            else:
                current = current.replace(month=current.month - 1)
            checked += 1
        print(f"[{prefix}] 找到 {len(indexes)} 个索引：{indexes}")
        return indexes

    def get_date_info_from_index(self, index_name):
        # 从索引名末尾提取 date_info，如 us_st_detail_month_2026_02 → 2026-02
        parts = index_name.rsplit('_', 2)
        return f"{parts[-2]}-{parts[-1]}"

    def build_profit_update_df(self, df):
        # 组装 profit_rate_extra 结构体，选出 update 所需列
        return df.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")
                )
            )
        ).select("asin", "profit_key", "profit_rate_extra")

    def write_profit_update(self, df, index_name):
        write_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.write.operation": "update",
            "es.batch.write.abort.on.failure": "false",
            "es.update.retry.on.conflict": "3",
            "es.batch.write.refresh": "false",
            "es.batch.size.entries": "2000",
            "es.batch.write.concurrency": "5",
            "es.batch.write.retry.count": "3",
            "es.batch.write.retry.wait": "60s",
            "es.nodes.wan.only": "false"
        }
        print(f"\n{'='*60}")
        print(f"开始更新索引利润率字段：{index_name}")
        print(f"{'='*60}")
        try:
            df.repartition(10).write.format("org.elasticsearch.spark.sql") \
                .options(**write_options) \
                .mode("append") \
                .save()
            print(f"索引 {index_name} 利润率字段更新完毕！")
        except Exception as e:
            print(f"更新索引 {index_name} 失败: {str(e)}")
            CommonUtil.send_wx_msg(['chenyuanjie'], '\u26A0 ES利润率更新失败', f'失败索引：{index_name}')

    def update_by_hive_asin(self, index_name, df_hive_asin):
        # Hive ASIN inner join 利润率增量数据，只更新有变化的 asin，减少 ES 写入量
        if not EsUtils.exist_index(index_name, self.es_client):
            print(f"索引 {index_name} 不存在，跳过更新")
            return
        df_update = df_hive_asin.join(self.df_asin_profit_rate, on='asin', how='inner')
        df_update = self.build_profit_update_df(df_update)
        self.write_profit_update(df_update, index_name)

    def update_by_es_asin_price(self, index_name):
        # 从 ES 读取 asin+price，按 asin+price 关联利润率后写回，确保 profit_key 匹配且不存在 document_missing
        if not EsUtils.exist_index(index_name, self.es_client):
            print(f"索引 {index_name} 不存在，跳过更新")
            return
        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": "true",
            "es.read.field.include": "asin,price"
        }
        df_es = self.spark.read.format("org.elasticsearch.spark.sql") \
            .options(**es_read_options) \
            .load(index_name) \
            .withColumn('price', F.round(F.col('price'), 2)) \
            .repartition(40, 'asin')
        df_update = df_es.join(self.df_asin_profit_rate, on=['asin', 'price'], how='inner')
        df_update = self.build_profit_update_df(df_update)
        self.write_profit_update(df_update, 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！！！")
