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


class EsAsinProfitRate(object):

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

        # ES相关配置
        self.es_client = EsUtils.get_es_client()
        self.es_index = f"{self.site_name}_profit_rate_extra"
        self.es_body = self.get_es_body()
        self.es_options = self.get_es_options(self.es_index)
        self.profit_rate_policy = f"{self.site_name}_profit_rate_policy"
        self.user_mask_asin_policy = "user_mask_asin_policy"
        self.user_mask_category_policy = "user_mask_category_policy"
        self.pipeline_id = f"{self.site_name}_user_mask_and_profit_rate_pipeline"

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

    @staticmethod
    def get_es_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"
                    }
                }
            }
        }

    @staticmethod
    def get_es_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_data()
        self.es_save()
        self.create_enrich_policy()
        self.create_enrich_pipeline()

    def read_data(self):
        sql = f"""
        select asin, price, ocean_profit, air_profit, package_length, package_width, package_height, weight 
        from dim_asin_profit_rate_info where site_name = '{self.site_name}'
        """
        self.df_asin_profit_rate = self.spark.sql(sqlQuery=sql).repartition(40, 'asin')

        sql = f"""
        select asin, package_length, package_width, package_height, weight 
        from dim_keepa_asin_info where site_name = '{self.site_name}'
        """
        self.df_keepa_asin = self.spark.sql(sqlQuery=sql).repartition(40, 'asin')

        # 因为 dim_asin_profit_rate_info 存在重复计算利润率的情况，保留与keepa最新数据所对应的数据行
        self.df_asin_profit_rate = self.df_asin_profit_rate.join(
            self.df_keepa_asin, on=['asin', 'package_length', 'package_width', 'package_height', 'weight'], how='inner'
        ).select(
            'asin', 'price', 'ocean_profit', 'air_profit'
        ).withColumn(
            'profit_key', F.concat_ws("_", F.col("asin"), F.col("price"))
        ).cache()

    def es_save(self):
        print(f"创建富集索引：{self.es_index}！")
        EsUtils.create_index(self.es_index, self.es_client, self.es_body)
        try:
            self.df_asin_profit_rate.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}')

    def create_enrich_policy(self):
        # self.es_client.ingest.delete_pipeline(id=self.pipeline_id)
        # self.es_client.enrich.delete_policy(name=self.policy_name)
        # print(f"创建富集策略：{self.policy_name}！")
        # policy_body = {
        #     "match": {
        #         "indices": f"{self.es_index}",
        #         "match_field": "profit_key",
        #         "enrich_fields": ["ocean_profit", "air_profit"]
        #     }
        # }
        # self.es_client.enrich.put_policy(name=self.policy_name, body=policy_body)

        print(f"刷新富集策略：{self.profit_rate_policy}！")
        self.es_client.enrich.execute_policy(self.profit_rate_policy, request_timeout=1800)

    def create_enrich_pipeline(self):
        print(f"创建富集管道：{self.pipeline_id}！")
        pipeline_body = {
            "description": "asin profit_rate and user_mask pipeline",
            "processors": [
                {
                    "enrich": {
                        "policy_name": self.profit_rate_policy,
                        "field": "profit_key",
                        "target_field": "profit_rate_extra",
                        "max_matches": 1,
                        "ignore_missing": True
                    },
                },
                {
                    "enrich": {
                        "policy_name": f"{self.user_mask_asin_policy}",
                        "field": "asin",
                        "target_field": "policy_add_1",
                        "max_matches": 1,
                        "ignore_missing": True
                    },
                },
                {
                    "enrich": {
                        "policy_name": f"{self.user_mask_category_policy}",
                        "field": "category_id",
                        "target_field": "policy_add_2",
                        "max_matches": 1,
                        "ignore_missing": True
                    },
                },
                {
                    "set": {
                        "field": "usr_mask_type",
                        "value": "{{policy_add_1.usr_mask_type}}",
                        "ignore_empty_value": True
                    }
                },
                {
                    "set": {
                        "field": "usr_mask_progress",
                        "value": "{{policy_add_1.usr_mask_progress}}",
                        "ignore_empty_value": True
                    }
                },
                {
                    "set": {
                        "field": "package_quantity",
                        "value": "{{policy_add_1.package_quantity}}",
                        "ignore_empty_value": True
                    }
                },
                {
                    "set": {
                        "field": "usr_mask_type",
                        "value": "{{policy_add_2.usr_mask_type}}",
                        "ignore_empty_value": True
                    }
                },
                {
                    "remove": {
                        "field": "policy_add_1",
                        "ignore_missing": True
                    }
                },
                {
                    "remove": {
                        "field": "policy_add_2",
                        "ignore_missing": True
                    }
                },
                {
                    "convert": {
                        "field": "package_quantity",
                        "type": "integer",
                        "ignore_missing": True
                    }
                }
            ]
        }

        self.es_client.ingest.put_pipeline(id=self.pipeline_id, body=pipeline_body)

        # 刷新ES数据，使pipeline生效
        # body = {
        #     "query": {
        #         "bool": {
        #             "must_not": {
        #                 "exists": {
        #                     "field": "profit_rate_extra"
        #                 }
        #             }
        #         }
        #     }
        # }
        body = {
        }
        self.es_client.update_by_query(
            index="us_st_detail_month_2025_11",
            body=body,
            pipeline=self.pipeline_id,
            refresh=True,
            wait_for_completion=False,
            request_timeout=600
        )
        pass


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