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 EsAiAsinAll(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"{site_name}_ai_analyze_extra"
        self.es_body = self.get_es_body()
        self.es_options = self.get_es_options(self.es_index)
        self.policy_name = f"{site_name}_ai_analyze_policy"
        self.pipeline_id = f"{site_name}_ai_analyze_pipeline"

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

    @staticmethod
    def get_es_body():
        return {
            "settings": {
                "number_of_shards": "3",
                "number_of_replicas": "1"
            },
            "mappings": {
                "properties": {
                    "asin": {
                        "type": "keyword"
                    },
                    "is_stable_flag": {
                        "type": "short"
                    },
                    "is_periodic_flag": {
                        "type": "short"
                    },
                    "is_ascending_flag": {
                        "type": "short"
                    },
                    "max_bought_month_arr": {
                        "type": "integer"
                    }
                }
            }
        }

    @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": "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"
        }

    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, 
            is_stable_flag, 
            is_periodic_flag, 
            is_ascending_flag, 
            max_month_last_12_month as max_bought_month_arr
        from dwt_ai_asin_all 
        where site_name = '{self.site_name}'
        """
        self.df_asin_detail = self.spark.sql(sqlQuery=sql).repartition(40, 'asin').withColumn(
            "max_bought_month_arr", F.split(F.col("max_bought_month_arr"), ",")
        ).withColumn(
            "max_bought_month_arr", F.expr("transform(max_bought_month_arr, x -> cast(x as int))")
        ).cache()
        print("ASIN信息库数据如下：")
        self.df_asin_detail.show(10, True)

    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_detail.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):
        # print(f"创建富集策略：{self.policy_name}！")
        # policy_body = {
        #     "match": {
        #         "indices": f"{self.es_index}",
        #         "match_field": "asin",
        #         "enrich_fields": ["is_stable_flag", "is_periodic_flag", "is_ascending_flag", "max_bought_month_arr"]
        #     }
        # }
        # self.es_client.enrich.put_policy(name=self.policy_name, body=policy_body)

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

    def create_enrich_pipeline(self):
        print(f"创建富集管道：{self.pipeline_id}！")
        pipeline_body = {
            "description": "ai asin analyze pipeline",
            "processors": [
                {
                    "enrich": {
                        "policy_name": self.policy_name,
                        "field": "asin",
                        "target_field": "last_year_extra",
                        "max_matches": 1,
                        "ignore_missing": True
                    },
                }
            ]
        }
        self.es_client.ingest.put_pipeline(id=self.pipeline_id, body=pipeline_body)

        pass


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