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.df_asin_profit_rate = self.spark.sql(f"select 1+1;")
        self.df_keepa_add = self.spark.sql(f"select 1+1;")
        self.df_cate_flag = self.spark.sql(f"select 1+1;")

    def run(self):
        self.read_profit_rate_add()
        self.read_keepa_add()
        # 利润率主索引写入Doris（已从ES迁移至Doris selection.{site_name}_asin_latest_profit_rate）
        self.save_profit_rate_to_doris()
        # 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)
            self.update_index_combined(index_name, self.get_month_last_day(date_info))
        # 年度信息库
        year_max_date_info = self.spark.sql(f"""
            select max(date_info) from dwt_ai_asin_year where site_name='{self.site_name}'
        """).collect()[0][0]
        if year_max_date_info is None:
            print(f"dwt_ai_asin_year 无数据，跳过年度索引更新")
        else:
            self.update_index_combined(
                f"{self.site_name}_ai_asin_analyze_detail_last365_day",
                self.get_month_last_day(str(year_max_date_info))
            )
        # ai_asin_analyze_detail 近3个月
        for index_name in self.get_recent_indexes("ai_asin_analyze_detail"):
            date_info = self.get_date_info_from_index(index_name)
            self.update_index_combined(index_name, self.get_month_last_day(date_info))
        # 30天流量选品（base_date=None → 以每个 asin 的 asin_crawl_date 为基准）
        self.read_cate_flag()
        self.update_index_combined(f"{self.site_name}_flow_asin_30day")
        # 释放缓存
        self.df_asin_profit_rate.unpersist()
        self.df_keepa_add.unpersist()
        self.df_cate_flag.unpersist()

    # ------------------------------------------------------------------ #
    # 数据读取
    # ------------------------------------------------------------------ #
    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和asin_price，用profit_key关联
        df_crawl_date = DorisHelper.spark_import_with_flight(
            session=self.spark,
            table_identifier=f"selection.{self.site_name}_asin_latest_detail",
            read_fields="asin,asin_crawl_date,asin_price"
        ).withColumn(
            "asin_crawl_date", F.substring(F.col("asin_crawl_date"), 1, 10)
        ).withColumn(
            "profit_key", F.concat_ws("_", F.col("asin"), F.round(F.col("asin_price"), 2))
        ).select('profit_key', 'asin_crawl_date') \
         .repartition(40, 'profit_key')

        self.df_asin_profit_rate = self.df_asin_profit_rate.join(
            df_crawl_date, on='profit_key', 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 read_keepa_add(self):
        # 读取keepa增量数据
        sql = f"""
        select asin, tracking_since from dim_keepa_asin_info
        where site_name = '{self.site_name}' and updated_time >= '{self.last_date_info}'
        """
        self.df_keepa_add = self.spark.sql(sqlQuery=sql).repartition(40, 'asin').cache()
        print(f"增量keepa数据如下：")
        self.df_keepa_add.show(10, False)

    def read_cate_flag(self):
        # 读取30day流量选品分类标记数据（用于更新30day索引额外字段）
        sql = f"""
        select asin,
               asin_cate_flag  as asin_source_flag,
               bsr_latest_date as bsr_last_seen_at,
               bsr_30day_count as bsr_seen_count_30d,
               nsr_latest_date as nsr_last_seen_at,
               nsr_30day_count as nsr_seen_count_30d
        from dwd_asin_cate_flag
        where site_name = '{self.site_name}' and date_type = '30day'
        """
        self.df_cate_flag = self.spark.sql(sqlQuery=sql).repartition(40, 'asin').cache()
        print(f"cate_flag数据量：{self.df_cate_flag.count()}")
        self.df_cate_flag.show(10, False)

    # ------------------------------------------------------------------ #
    # 利润率主索引写入Doris
    # ------------------------------------------------------------------ #
    def save_profit_rate_to_doris(self):
        doris_table = f"{self.site_name}_asin_latest_profit_rate"
        print(f"\n{'='*60}")
        print(f"开始更新Doris利润率表：selection.{doris_table}")
        print(f"{'='*60}")
        table_columns = 'asin,price,ocean_profit,air_profit,asin_crawl_date,update_time'
        DorisHelper.spark_export_with_columns(
            df_save=self.df_asin_profit_rate.select(
                'asin', 'price', 'ocean_profit', 'air_profit',
                F.to_date(F.col('asin_crawl_date'), 'yyyy-MM-dd').alias('asin_crawl_date'),
                F.to_date(F.col('update_time'), 'yyyy-MM-dd').alias('update_time')
            ),
            db_name='selection',
            table_name=doris_table,
            table_columns=table_columns,
            use_type='selection',
        )
        print(f"Doris selection.{doris_table} 更新完毕！")

    # ------------------------------------------------------------------ #
    # 工具方法
    # ------------------------------------------------------------------ #
    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 get_month_last_day(self, year_month_str):
        # year_month_str: "yyyy-MM"，返回该月最后一天 "yyyy-MM-dd"
        year, month = map(int, year_month_str.split('-'))
        if month == 12:
            last_day = datetime(year + 1, 1, 1) - timedelta(days=1)
        else:
            last_day = datetime(year, month + 1, 1) - timedelta(days=1)
        return last_day.strftime('%Y-%m-%d')

    def get_tracking_interval_dict(self, base_date):
        base = datetime.strptime(base_date, '%Y-%m-%d')
        return {
            "one_month":         (base - timedelta(days=30)).strftime('%Y-%m-%d'),
            "three_month":       (base - timedelta(days=90)).strftime('%Y-%m-%d'),
            "six_month":         (base - timedelta(days=180)).strftime('%Y-%m-%d'),
            "twelve_month":      (base - timedelta(days=360)).strftime('%Y-%m-%d'),
            "twenty_four_month": (base - timedelta(days=720)).strftime('%Y-%m-%d'),
            "thirty_six_month":  (base - timedelta(days=1080)).strftime('%Y-%m-%d'),
        }

    def get_tracking_since_type_expr(self, iv):
        return f"""
            CASE WHEN tracking_since >= '{iv['one_month']}' THEN 1
                 WHEN tracking_since >= '{iv['three_month']}' AND tracking_since < '{iv['one_month']}' THEN 2
                 WHEN tracking_since >= '{iv['six_month']}' AND tracking_since < '{iv['three_month']}' THEN 3
                 WHEN tracking_since >= '{iv['twelve_month']}' AND tracking_since < '{iv['six_month']}' THEN 4
                 WHEN tracking_since >= '{iv['twenty_four_month']}' AND tracking_since < '{iv['twelve_month']}' THEN 5
                 WHEN tracking_since >= '{iv['thirty_six_month']}' AND tracking_since < '{iv['twenty_four_month']}' THEN 6
                 WHEN tracking_since < '{iv['thirty_six_month']}' THEN 7
                 ELSE 0 END
        """

    # ------------------------------------------------------------------ #
    # ES 读写配置
    # ------------------------------------------------------------------ #
    def get_es_read_options(self, include_crawl_date=False):
        # es.mapping.date.rich=false：date字段以原始字符串返回，避免转为timestamp带来的精度问题
        # profit_rate_extra 不从ES读取：月度索引所有文档初始值为{}，ES-Hadoop推断为空struct，
        # 写回时F.when/otherwise两分支类型不一致导致struct字段被丢弃，序列化为{}
        # 改为直接从Hive新值构建struct写入，彻底绕开空struct的schema推断问题
        fields = "asin,profit_key,tracking_since,tracking_since_type"
        if include_crawl_date:
            fields += ",asin_crawl_date"
        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.nodes.wan.only": "false",
            "es.mapping.date.rich": "false",
            "es.read.field.include": fields,
            "es.scroll.size": "5000"
        }

    def write_combined_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": "10",
            "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_index_combined(self, index_name, base_date=None):
        """
        base_date: 固定基准日期（yyyy-MM-dd），用于月度/年度索引 tracking_since_type 计算
                   None 表示 30day 索引，以每个 asin 的 asin_crawl_date 为基准动态计算
        """
        if not EsUtils.exist_index(index_name, self.es_client):
            print(f"索引 {index_name} 不存在，跳过更新")
            return
        print(f"\n{'='*60}")
        print(f"开始处理索引：{index_name}，base_date={base_date or 'asin_crawl_date(动态)'}")
        print(f"{'='*60}")

        # 1. 读取ES现有数据，date字段以原始字符串返回
        # profit_rate_extra 不读取（见 get_es_read_options 注释）
        # asin_crawl_date 仅 30day 动态模式需要，按需读取
        # cache：30day索引步骤8需复用 df_es.select('asin')，避免二次全量读取
        df_es = self.spark.read.format("org.elasticsearch.spark.sql") \
            .options(**self.get_es_read_options(include_crawl_date=(base_date is None))) \
            .load(index_name) \
            .withColumn('tracking_since_type', F.col('tracking_since_type').cast('int')) \
            .repartition(40, 'asin') \
            .cache()

        # 2. 准备利润率增量（直接用 profit_key 关联，无需读取 ES price 字段）
        df_profit = self.df_asin_profit_rate.select(
            'profit_key',
            F.col('ocean_profit').cast('float').alias('new_ocean_profit'),
            F.col('air_profit').cast('float').alias('new_air_profit')
        )

        # 3. 准备keepa增量（重命名避免列冲突）
        df_keepa = self.df_keepa_add.select(
            'asin',
            F.col('tracking_since').alias('kp_tracking_since')
        )

        # 4. left join 两个增量数据集
        # profit_key 格式统一（asin_price，已 round），ES 与 Hive 两侧直接对齐，无需借助 price 列
        df = df_es \
            .join(df_profit, on='profit_key', how='left') \
            .join(df_keepa, on='asin', how='left')

        # 5. 过滤：至少一方有增量数据，减少无效写入
        df = df.filter(F.col('new_ocean_profit').isNotNull() | F.col('new_air_profit').isNotNull() | F.col('kp_tracking_since').isNotNull())

        # 6. keepa字段合并：有新数据用新的，否则保留ES旧值
        # 将 kp_tracking_since（keepa原始分钟值）转为日期字符串
        df = df.withColumn('tracking_since',
            F.when(
                F.col('kp_tracking_since').isNotNull(),
                F.date_format(
                    F.from_unixtime((F.col('kp_tracking_since') + F.lit(21564000)) * 60),
                    'yyyy-MM-dd'
                )
            ).otherwise(F.col('tracking_since'))
        )

        # 计算 tracking_since_type
        if base_date is not None:
            # 固定基准：月度/年度索引
            iv = self.get_tracking_interval_dict(base_date)
            df = df.withColumn('tracking_since_type',
                F.when(F.col('kp_tracking_since').isNotNull(), F.expr(self.get_tracking_since_type_expr(iv)))
                 .otherwise(F.col('tracking_since_type'))
            )
        else:
            # 动态基准：30day索引，以每个 asin 的 asin_crawl_date 为基准
            # asin_crawl_date 格式为 "yyyy-MM-dd HH:mm:ss"，截取前10位转为日期字符串
            # 将 tracking_since 字符串显式转为 DateType，避免与 date_sub 返回值类型不一致
            df = df.withColumn('asin_crawl_date', F.substring(F.col('asin_crawl_date'), 1, 10))
            ts_date = F.to_date(F.col('tracking_since'), 'yyyy-MM-dd')
            df = df.withColumn('tracking_since_type',
                F.when(F.col('kp_tracking_since').isNotNull(),
                    F.when(ts_date >= F.date_sub(F.col('asin_crawl_date'), 30), F.lit(1))
                     .when(ts_date >= F.date_sub(F.col('asin_crawl_date'), 90), F.lit(2))
                     .when(ts_date >= F.date_sub(F.col('asin_crawl_date'), 180), F.lit(3))
                     .when(ts_date >= F.date_sub(F.col('asin_crawl_date'), 360), F.lit(4))
                     .when(ts_date >= F.date_sub(F.col('asin_crawl_date'), 720), F.lit(5))
                     .when(ts_date >= F.date_sub(F.col('asin_crawl_date'), 1080), F.lit(6))
                     .when(ts_date < F.date_sub(F.col('asin_crawl_date'), 1080), F.lit(7))
                     .otherwise(F.lit(0))
                ).otherwise(F.col('tracking_since_type'))
            )

        # 7. 写回索引（tracking_since_type 写回 ES short 类型，需显式 cast）
        # 拆为两次写入：
        #   7a. 有利润率数据的行 → profit_rate_extra struct 从新值直接构建，彻底绕开空struct schema推断问题
        #   7b. 仅有 keepa 数据的行 → 不写 profit_rate_extra，避免 null 覆盖已有利润率值
        df = df.cache()

        df_profit_update = df.filter(
            F.col('new_ocean_profit').isNotNull() | F.col('new_air_profit').isNotNull()
        ).select(
            'asin', 'profit_key',
            F.struct(
                F.col('new_ocean_profit').cast('float').alias('ocean_profit'),
                F.col('new_air_profit').cast('float').alias('air_profit')
            ).alias('profit_rate_extra'),
            'tracking_since',
            F.col('tracking_since_type').cast('short')
        )
        self.write_combined_update(df_profit_update, index_name)

        df_keepa_update = df.filter(
            F.col('new_ocean_profit').isNull() & F.col('new_air_profit').isNull()
        ).select(
            'asin', 'profit_key', 'tracking_since',
            F.col('tracking_since_type').cast('short')
        )
        self.write_combined_update(df_keepa_update, index_name)

        df.unpersist()

        # 8. 30day 索引额外更新 cate_flag 相关字段（partial update，不影响其他字段）
        # inner join df_es 确保只更新索引中已存在的 asin，避免 doc missing 报错
        # asin_source_flag：Hive 存为字符串，ES mapping 为 integer[]，需转为 int 数组
        if base_date is None:
            print(f"[30day] 开始更新 cate_flag 字段：asin_source_flag / bsr / nsr")
            df_cate_update = self.df_cate_flag.join(
                df_es.select('asin'), on='asin', how='inner'
            ).select(
                'asin',
                F.transform(F.split(F.col('asin_source_flag'), ','), lambda x: x.cast('int')).alias('asin_source_flag'),
                F.coalesce(F.col('bsr_last_seen_at'), F.lit('1970-01-01')).alias('bsr_last_seen_at'),
                F.coalesce(F.col('bsr_seen_count_30d').cast('int'), F.lit(0)).alias('bsr_seen_count_30d'),
                F.coalesce(F.col('nsr_last_seen_at'), F.lit('1970-01-01')).alias('nsr_last_seen_at'),
                F.coalesce(F.col('nsr_seen_count_30d').cast('int'), F.lit(0)).alias('nsr_seen_count_30d')
            )
            self.write_combined_update(df_cate_update, index_name)

        df_es.unpersist()


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！！！")
