# !/usr/bin/env python
# -*- coding: utf-8 -*-
import copy
import json,re,sys,os
from elasticsearch import Elasticsearch, helpers  #  elasticsearch==7.0.0； pandas==2.2.3；  numpy==2.0.2
import jieba
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, as_completed

# 将 非AI资源 json 数据保存进es数据库


class EsHelper:
    def __init__(self, hosts=None):
        """初始化ES连接
        Args:
            hosts: ES服务器地址列表，默认为None使用已配置的地址
        """
        if hosts is None:
            self.hosts = [{'host': '116.63.110.220', 'port': 9200}]
            self.hosts = [{'host': 'localhost', 'port': 9200}]
            self.hosts = ['http://localhost:9200']  # http://116.63.110.220:9200
        self.es = Elasticsearch(hosts=self.hosts, timeout=1000)
        print("self.es.ping(): ", self.es.ping())

    def create_index(self, index_name):
        """创建索引
        Args:
            index_name: es索引名称
            settings: 索引设置
            mappings: 索引映射
            ['tableName','bookCategory', 'typeCode', 'raysAppId', 'createTime',  'resourceName', 'subjectCategory', 'studentSemester',
            'audience', 'studentGrade',   'resourceDescribe', 'resourceContent', 'subjectORlabel',
             'gradeDeep',  'resourceFrom',
            'workType', 'schoolStage', 'bookVersion', 'purpose', 'authorName',  'textBookEdition', 'message', 'bookPrice'
            ]
        Returns:
            创建结果
        """
        body = {}
        settings = {
            "index": {
                "number_of_shards": 1,  # 定义索引的分片数量 方便在不同的节点部署
                "number_of_replicas": 0,  # 定义索引中每个主分片的副本数量
            }}

        mappings = {
            "dynamic": True,
            "properties": {
                "tableName": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "bookCategory": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "typeCode": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "raysAppId": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "createTime": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "resourceName": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "subjectCategory": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "studentSemester": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "audience": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "studentGrade": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },  # 10
                "resourceDescribe": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "resourceContent": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "subjectORlabel": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "gradeDeep": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "resourceFrom": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "workType": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "schoolStage": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "bookVersion": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "purpose": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "authorName": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },  # 20
                "textBookEdition": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "message": {
                    "type": "text",
                    "fields": {
                        "keyword": {
                            "type": "keyword",
                            "ignore_above": 256
                        }
                    }
                },
                "bookPrice": {
                    "type": "float",

                },



            }
        }
        if settings:
            body['settings'] = settings
        if mappings:
            body['mappings'] = mappings
        res = self.es.indices.create(index=index_name, body=body, timeout='100s')
        print(f"create index {index_name} success.")
        return res

    def delete_index(self, index_name):
        """删除索引
        Args:
            index_name: 索引名称
        Returns:
            删除结果
        """
        res = self.es.indices.delete(index=index_name)
        print(f"删除索引 {index_name} success.")
        return res

    def index_exists(self, index_name):
        """判断索引是否存在

        Args:
            index_name: 索引名称
        Returns:
            bool: 是否存在
        """
        return self.es.indices.exists(index=index_name)

    def read_dict(self, file):
        # 1、读入中朝大辞典 ocr解析 并 gpt4o清洗后的数据
        with open(file, 'r', encoding='utf-8') as f:
            data_list = json.loads(f.read())

            res = []
            for item in data_list:
                temp = {}
                for k, v in item.items():
                    if not v:
                        pass
                    else:
                        temp[k] = v
                res.append(copy.deepcopy(temp))

            print("len(data_list): ", len(data_list))
            print("len(res): ", len(res))

        return res

    def insert_doc(self, index_name, doc_body, doc_id=None):
        """插入文档
        Args:
            index_name: 索引名称
            doc_body: 文档内容
            doc_id: 文档ID,可选
        Returns:
            插入结果
        """
        return self.es.index(index=index_name, body=doc_body, id=doc_id)

    def insert_doc_batch(self, index_name, doc_batch: list):
        # 批量插入es数据库
        actions = []
        for idx, body in enumerate(doc_batch):
            if 'id' in body:
                idx = body['id']
            temp = {"_index": index_name, "_id": idx}
            temp['_source'] = body
            actions.append(temp)

        # 使用helpers.bulk方法批量插入文档
        helpers.bulk(self.es, actions)

    def get_doc(self, index_name, doc_id):
        """获取文档
        Args:
            index_name: 索引名称
            doc_id: 文档ID
        Returns:
            文档内容
        """
        return self.es.get(index=index_name, id=doc_id)

    def update_doc(self, index_name, doc_id, doc_body):
        """更新文档

        Args:
            index_name: 索引名称
            doc_id: 文档ID
            doc_body: 更新的内容
        Returns:
            更新结果
        """
        return self.es.update(index=index_name, id=doc_id, body={'doc': doc_body})

    def delete_doc(self, index_name, doc_id):
        """删除文档

        Args:
            index_name: 索引名称
            doc_id: 文档ID
        Returns:
            删除结果
        """
        return self.es.delete(index=index_name, id=doc_id)

    def search(self, index_name, query_body):
        """搜索文档
        """
        return self.es.search(index=index_name, body=query_body)


class ESsearch():
    # 搜索es数据库
    def __init__(self, hosts=['http://localhost:9200']):
        self.es = Elasticsearch(hosts=hosts, timeout=30)
        self.index_name = 'no_ai_source'

        self.size_study_tool = 5  # 伴学工具(这个表没有使用)
        self.size_third_party_own_resource = 50  # 第三方自有资源评级详表  100
        self.size_brand_resource = 35    # 品牌资源（爱奇艺+慕课）
        self.size_xiaorui_information = 10   # 小睿资讯服务
        self.size_evaluation_database = 5   # 测评库资源(这个表没有使用)

        # 【学段："schoolStage" 小学】；【学科： "subjectCategory", 语文】；【年级： "studentGrade" 一年级】。


    def execute_query(self, query=None,  size=None, offset=0):
        # 执行 非AI资源包 es查询语句
        if query is None:
            query = {"bool": {
                                  "should": [{"match_phrase": {"message": "K12"}}]
                              }
                    }


        query_body = {"from": offset,
                          "query": query
                          }
        if size:
            query_body['size'] = size

        res = self.es.search(index=self.index_name, body=query_body)

        res_hits = []
        if 'hits' in res and 'hits' in res['hits']:
            hits = res['hits']['hits']
            for item in hits:
                if "_source" in item and item['_source']:
                    res_hits.append(item['_source'])

        for item in res_hits:
            for k, v in item.items():
                if  isinstance(v, str):
                    v = re.sub('"', " ", v)  # 英文双引号转中文双引号
                    v = re.sub("'", " ", v) # “
                    v = re.sub("“", " ", v) #
                    v = re.sub("”", " ", v) # ““”  中文双引号
                    item[k] = v

        return res_hits

    def search_study_tool(self, input):
        # 1、es搜索 伴学工具    全部是K12的资源
        query = {"bool": {"should": [],
                          "must":[]
                          }
                 }
        must_query = [
                      {"match_phrase": {"tableName": "伴学工具"}},
            {"match_phrase": {"bookCategory": input['bookCategory']}}
                      ]
        should_query = []
        for k, v in input.items():
            if not v:
                continue
            if k in [ "subjectCategory"] and v:  # 学科 语数外
                temp = {"match_phrase": {"message": v}}
                must_query.append(temp)

            else:
                # temp = {"match": {"message": v}}
                temp = {"match": {"resourceName": v}}
                should_query.append(temp)

        query['bool']['should'] = should_query
        query['bool']['must'] = must_query

        return query

    def search_third_party_own_resource(self, input):
        # 2、es搜索 第三方自有资源评级详表
        # print("2、es搜索 第三方自有资源评级详表")
        # print(input)
        """
        input = {
            "subjectCategory": "语文",
            "schoolStage": "小学",
            "studentGrade": "二年级",
            "purpose": "基础学习",
            "bookVersion": "人教版",
            "keyword": "现代汉语；古代汉语；中国文学常识；外国文学常识；阅读；写作；教学"
        }
        """

        query = {"bool": {"should": [],
                          "must": []
                          }
                 }
        must_query = [
                      {"match_phrase": {"tableName": "第三方自有资源评级详表"}},
            {"match_phrase": {"bookCategory": input['bookCategory']}}
                      ]
        should_query = []
        for k, v in input.items():
            if not v:
                continue
            if input['subjectCategory']=="K12" and k in ["schoolStage", "studentGrade", "subjectCategory", "bookVersion"] and v:
                v = v.strip().split("；")
                v = list(set(v))
                for v0 in v:
                    temp = {"match_phrase": {"message": v0}}
                    must_query.append(temp)

            else:
                # temp = {"match": {"message": v}}
                temp = {"match": {"resourceName": v}}
                should_query.append(temp)

        query['bool']['should'] = should_query
        query['bool']['must'] = must_query

        return query

    def search_brand_resource(self, input):
        # 3、es搜索 品牌资源（爱奇艺+慕课）
        query = {"bool": {"should": [],
                          "must": []
                          }
                 }
        must_query = [
            {"match_phrase": {"tableName": "品牌资源（爱奇艺+慕课）"}},
        {"match_phrase": {"bookCategory": input['bookCategory']}}
        ]
        should_query = []
        for k, v in input.items():
            if not v:
                continue
            if input['subjectCategory']=="K12" and k in ["schoolStage", "subjectCategory",  "studentGrade"] and v:
                v = v.strip().split("；")
                v = list(set(v))
                for v0 in v:
                    temp = {"match_phrase": {"message": v0}}
                    must_query.append(temp)

            else:
                # temp = {"match": {"message": v}}
                temp = {"match": {"resourceName": v}}
                should_query.append(temp)

        query['bool']['should'] = should_query
        query['bool']['must'] = must_query

        return query

    def search_xiaorui_information(self, input):
        # 4、es搜索 小睿资讯服务
        query = {"bool": {"should": [],
                          "must": []
                          }
                 }
        must_query = [
            {"match_phrase": {"tableName": "小睿资讯服务"}},
        {"match_phrase": {"bookCategory": input['bookCategory']}}
        ]
        should_query = []
        for k, v in input.items():
            if not v:
                continue
            if input['bookCategory']=="K12" and k in ["schoolStage", "subjectCategory",  "studentGrade"] and v:
                v = v.strip().split("；")
                v = list(set(v))
                for v0 in v:
                    temp = {"match_phrase": {"message": v0}}
                    must_query.append(temp)

            else:
                # temp = {"match": {"message": v}}
                temp = {"match": {"resourceName": v}}
                should_query.append(temp)

        query['bool']['should'] = should_query
        query['bool']['must'] = must_query

        return query

    def search_evaluation_database(self, input):
        # 5、es搜索 测评库资源
        query = {"bool": {"should": [],
                          "must": []
                          }
                 }
        must_query = [
            {"match_phrase": {"tableName": "测评库资源"}},
            {"match_phrase": {"bookCategory": input['bookCategory']}}
        ]
        should_query = []
        for k, v in input.items():
            if not v:
                continue
            if input['bookCategory']=="K12" and k in ["schoolStage", "subjectCategory",  "studentGrade"] and v:
                v = v.strip().split("；")
                v = list(set(v))
                for v0 in v:
                    if not v0:
                        continue
                    temp = {"match_phrase": {"message": v0}}
                    must_query.append(temp)

            else:
                # temp = {"match": {"message": v}}
                temp = {"match": {"resourceName": v}}
                should_query.append(temp)

        query['bool']['should'] = should_query
        query['bool']['must'] = must_query

        return query

    def search_no_ai_source(self, input):
        query_study_tool = self.search_study_tool(input)
        query_third_party_own_resource = self.search_third_party_own_resource(input)
        query_brand_resource = self.search_brand_resource(input)
        query_xiaorui_information = self.search_xiaorui_information(input)
        query_evaluation_database = self.search_evaluation_database(input)

        # K12 图书
        if input['bookCategory'] == "K12":
            # 并行执行5个函数
            with ThreadPoolExecutor(max_workers=5) as executor:
                res1 = executor.submit(self.execute_query, query=query_study_tool, size=self.size_study_tool)
                res2 = executor.submit(self.execute_query, query=query_third_party_own_resource, size=self.size_third_party_own_resource)
                res3 = executor.submit(self.execute_query, query=query_brand_resource, size=self.size_brand_resource)
                res4 = executor.submit(self.execute_query, query=query_xiaorui_information, size=self.size_xiaorui_information)
                res5 = executor.submit(self.execute_query, query=query_evaluation_database, size=self.size_evaluation_database)

                # 等待所有任务完成并获取结果
                res1 = res1.result()
                res2 = res2.result()
                res3 = res3.result()
                res4 = res4.result()
                res5 = res5.result()

                res = res1 + res2 + res3 + res4 + res5

        # 童书 和 其他图书
        else:
            # 并行执行4个函数
            with ThreadPoolExecutor(max_workers=4) as executor:
                res2 = executor.submit(self.execute_query, query=query_third_party_own_resource, size=self.size_third_party_own_resource)
                res3 = executor.submit(self.execute_query, query=query_brand_resource, size=self.size_brand_resource)
                res4 = executor.submit(self.execute_query, query=query_xiaorui_information, size=self.size_xiaorui_information)
                res5 = executor.submit(self.execute_query, query=query_evaluation_database, size=self.size_evaluation_database)

                # 等待所有任务完成并获取结果
                res2 = res2.result()
                res3 = res3.result()
                res4 = res4.result()
                res5 = res5.result()

                res = res2 + res3 + res4 + res5

        return res

    def search(self, input):
        # 后端调用的搜索接口
        if input['bookCategory'] not in ["K12", "童书"]:
            input['bookCategory'] = "其他图书"

        if "tableName" not in input.keys() or not input['tableName']:
            res = self.search_no_ai_source(input)
        elif "tableName"  in input.keys() and input['tableName'] =="伴学工具":
            query_study_tool = self.search_study_tool(input)
            res = self.execute_query( query = query_study_tool, size = self.size_study_tool)
        elif "tableName" in input.keys() and input['tableName'] == "第三方自有资源评级详表":
            query_third_party_own_resource = self.search_third_party_own_resource(input)
            res = self.execute_query(query=query_third_party_own_resource, size=self.size_third_party_own_resource)
        elif "tableName" in input.keys() and input['tableName'] == "品牌资源（爱奇艺+慕课）":
            query_brand_resource = self.search_brand_resource(input)
            res = self.execute_query(query=query_brand_resource, size=self.size_brand_resource)
        elif "tableName" in input.keys() and input['tableName'] == "小睿资讯服务":
            query_xiaorui_information = self.search_xiaorui_information(input)
            res = self.execute_query(query=query_xiaorui_information, size=self.size_xiaorui_information)
        elif "tableName" in input.keys() and input['tableName'] == "测评库资源":
            query_evaluation_database = self.search_evaluation_database(input)
            res = self.execute_query(query=query_evaluation_database, size=self.size_evaluation_database)

        else:
            res = [{"error":"tableName must in 【伴学工具；第三方自有资源评级详表；品牌资源（爱奇艺+慕课）；小睿资讯服务；测评库资源】 "}]


        # 踢出为空的字段
        for item in res:
            for k, v in item.items():
                if not v:
                    item.pop(k)
        return res


if __name__ == "__main__":
    index_name = "no_ai_source"  # 非AI资源 数据库
    es = EsHelper()
    # # #es.delete_index( index_name) # 删除es索引
    # print(es.create_index("test_index", {"number_of_shards": 1, "number_of_replicas": 0}))
    # print(es.search("test_index", {"query": {"match_all": {}}}))

    # # # 1、创建新的es索引
    # if es.index_exists( index_name):
    #     es.delete_index(index_name)  # 删除索引
    # create_index_res = es.create_index(index_name)  # 建立一个新的es 索引
    # print('create_index_res: ', create_index_res)
    #
    # # # 2、es中插入数据
    # file = '/home/liuxin/work/AI_planner/data/非AI工具/商务资源SOP管理表es.json'
    # data_list = es.read_dict( file)  # 18169
    # es.insert_doc_batch( index_name, data_list)

    # res = es.get_doc(index_name, 18128)
    # print(res)
    #
    # # 3、查看es中有多少条数据
    # count = es.es.count(index_name)
    # print(f'{index_name} ：{count} ')  #


    ######  es 搜索 非AI资源数据
    es_search = ESsearch()
    input = {
"bookCategory": "童书",
"subjectCategory": "语文",
"keyword": "查询方便、词库齐全"

}
    hits = es_search.search(input)
    # hits = es_search.search_no_ai_source(input)

    hits = json.dumps(hits, ensure_ascii=False, indent=4)
    print(hits)



    print('\nfinished.')