#
# Copyright 2023 OpenSPG Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
# in compliance with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software distributed under the License
# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
# or implied.

import json
import re
from typing import List, Dict, Any
from collections import defaultdict

from knext.schema.model.schema_helper import SPGTypeName
from knext.builder.operator.spg_record import SPGRecord
from knext.builder.operator.builtin.auto_prompt import AutoPrompt


class OneKEPrompt(AutoPrompt):
    template_zh: str = ""
    template_en: str = ""

    def __init__(self, **kwargs):
        types_list = kwargs.get("types_list", [])
        language = kwargs.get("language", "zh")
        with_description = kwargs.get("language", False)
        split_num = kwargs.get("split_num", 4)
        super().__init__(types_list)
        if language == "zh":
            self.template = self.template_zh
        else:
            self.template = self.template_en
        self.with_description = with_description
        self.split_num = split_num

        self._init_render_variables()
        self._render()

        self.params = kwargs

    def build_prompt(self, variables: Dict[str, str]) -> List[str]:
        instructions = []
        for schema in self.schema_list:
            instructions.append(
                json.dumps(
                    {
                        "instruction": self.template,
                        "schema": schema,
                        "input": variables.get("input"),
                    },
                    ensure_ascii=False,
                )
            )
        return instructions

    def parse_response(self, response: str) -> List[SPGRecord]:
        raise NotImplementedError

    def _render(self):
        raise NotImplementedError

    def multischema_split_by_num(self, split_num, schemas: List[Any]):
        negative_length = max(len(schemas) // split_num, 1) * split_num
        total_schemas = []
        for i in range(0, negative_length, split_num):
            total_schemas.append(schemas[i : i + split_num])

        remain_len = max(1, split_num // 2)
        tmp_schemas = schemas[negative_length:]
        if len(schemas) - negative_length >= remain_len and len(tmp_schemas) > 0:
            total_schemas.append(tmp_schemas)
        elif len(tmp_schemas) > 0:
            total_schemas[-1].extend(tmp_schemas)
        return total_schemas


class OneKE_NERPrompt(OneKEPrompt):
    template_zh: str = (
        "你是专门进行实体抽取的专家。请从input中抽取出符合schema定义的实体，不存在的实体类型返回空列表。请按照JSON字符串的格式回答。"
    )
    template_en: str = "You are an expert in named entity recognition. Please extract entities that match the schema definition from the input. Return an empty list if the entity type does not exist. Please respond in the format of a JSON string."

    def __init__(
        self,
        entity_types: List[SPGTypeName],
        language: str = "zh",
        with_description: bool = False,
        split_num: int = 4,
    ):
        super().__init__(
            types_list=entity_types,
            language=language,
            with_description=with_description,
            split_num=split_num,
        )

    def parse_response(self, response: str) -> List[SPGRecord]:
        if isinstance(response, list) and len(response) > 0:
            response = response[0]
        try:
            ent_obj = json.loads(response)
        except json.decoder.JSONDecodeError:
            print("DeepKE_NERPrompt response JSONDecodeError error.")
            return []
        if type(ent_obj) != dict:
            print("DeepKE_NERPrompt response type error.")
            return []

        spg_records = []
        for type_zh, values in ent_obj.items():
            if type_zh not in self.spg_type_schema_info_zh:
                print(f"Unrecognized entity_type: {type_zh}")
                continue
            type_en, _ = self.spg_type_schema_info_zh[type_zh]
            spg_record = SPGRecord(type_en)
            for value in values:
                spg_record.upsert_properties({"id": value, "name": value})
            spg_records.append(spg_record)

    def _render(self):
        entity_list = []
        for spg_type in self.spg_types:
            entity_list.append(spg_type.name_zh)
        self.schema_list = self.multischema_split_by_num(self.split_num, entity_list)


class OneKE_SPOPrompt(OneKEPrompt):
    template_zh: str = (
        "你是专门进行SPO三元组抽取的专家。请从input中抽取出符合schema定义的spo关系三元组，不存在的关系返回空列表。请按照JSON字符串的格式回答。"
    )
    template_en: str = "You are an expert in spo(subject, predicate, object) triples extraction. Please extract SPO relationship triples that match the schema definition from the input. Return an empty list for relationships that do not exist. Please respond in the format of a JSON string."

    def __init__(
        self,
        spo_types: List[SPGTypeName],
        language: str = "zh",
        with_description: bool = False,
        split_num: int = 4,
    ):
        super().__init__(
            types_list=spo_types,
            language=language,
            with_description=with_description,
            split_num=split_num,
        )
        self.properties_mapper = {}
        self.relations_mapper = {}

    def parse_response(self, response: str) -> List[SPGRecord]:
        if isinstance(response, list) and len(response) > 0:
            response = response[0]
        try:
            re_obj = json.loads(response)
        except json.decoder.JSONDecodeError:
            print("DeepKE_REPrompt response JSONDecodeError error.")
            return []
        if type(re_obj) != dict:
            print("DeepKE_REPrompt response type error.")
            return []

        relation_dcir = defaultdict(list)
        for relation_zh, values in re_obj.items():
            if relation_zh not in self.property_info_zh[relation_zh]:
                print(f"Unrecognized relation: {relation_zh}")
                continue
            if values and isinstance(values, list):
                for value in values:
                    if (
                        type(value) != dict
                        or "subject" not in value
                        or "object" not in value
                    ):
                        print("DeepKE_REPrompt response type error.")
                        continue
                    s_zh, o_zh = value.get("subject", ""), value.get("object", "")
                    relation_dcir[relation_zh].append((s_zh, o_zh))

        spg_records = []
        for relation_zh, sub_obj_list in relation_dcir.items():
            sub_dict = defaultdict(list)
            for s_zh, o_zh in sub_obj_list:
                sub_dict[s_zh].append(o_zh)
            for s_zh, o_list in sub_dict.items():
                if s_zh in self.spg_type_schema_info_zh:
                    print(f"Unrecognized subject_type: {s_zh}")
                    continue
                object_value = ",".join(o_list)
                s_type_zh = self.properties_mapper.get(relation_zh, None)
                if s_type_zh is not None:
                    s_type_en, _ = self.spg_type_schema_info_zh[s_type_zh]
                    relation_en, _ = self.property_info_zh[relation_zh]
                    spg_record = SPGRecord(s_type_en).upsert_properties(
                        {"id": s_zh, "name": s_zh}
                    )
                    spg_record.upsert_property(relation_en, object_value)
                else:
                    s_type_zh, o_type_zh = self.relations_mapper.get(
                        relation_zh, [None, None]
                    )
                    if s_type_zh is None or o_type_zh is None:
                        print(f"Unrecognized relation: {relation_zh}")
                        continue
                    s_type_en, _ = self.spg_type_schema_info_zh[s_type_zh]
                    spg_record = SPGRecord(s_type_en).upsert_properties(
                        {"id": s_zh, "name": s_zh}
                    )
                    relation_en, _, object_type = self.relation_info_zh[s_type_zh][
                        relation_zh
                    ]
                    spg_record.upsert_relation(relation_en, object_type, object_value)
                spg_records.append(spg_record)
        return spg_records

    def _render(self):
        spo_list = []
        for spg_type in self.spg_types:
            type_en, _ = self.spg_type_schema_info_zh[spg_type]
            for v in spg_type.properties.values():
                spo_list.append(
                    {
                        "subject_type": spg_type.name_zh,
                        "predicate": v.name_zh,
                        "object_type": "文本",
                    }
                )
                self.properties_mapper[v.name_zh] = spg_type
            for v in spg_type.relations.values():
                _, _, object_type = self.relation_info_en[type_en][v.name]
                spo_list.append(
                    {
                        "subject_type": spg_type.name_zh,
                        "predicate": v.name_zh,
                        "object_type": object_type,
                    }
                )
                self.relations_mapper[v.name_zh] = [spg_type, object_type]
        self.schema_list = self.multischema_split_by_num(self.split_num, spo_list)


class OneKE_REPrompt(OneKE_SPOPrompt):
    template_zh: str = (
        "你是专门进行关系抽取的专家。请从input中抽取出符合schema定义的关系三元组，不存在的关系返回空列表。请按照JSON字符串的格式回答。"
    )
    template_en: str = "You are an expert in relationship extraction. Please extract relationship triples that match the schema definition from the input. Return an empty list for relationships that do not exist. Please respond in the format of a JSON string."

    def __init__(
        self,
        relation_types: List[SPGTypeName],
        language: str = "zh",
        with_description: bool = False,
        split_num: int = 4,
    ):
        super().__init__(relation_types, language, with_description, split_num)

    def _render(self):
        re_list = []
        for spg_type in self.spg_types:
            type_en, _ = self.spg_type_schema_info_zh[spg_type]
            for v in spg_type.properties.values():
                re_list.append(v.name_zh)
                self.properties_mapper[v.name_zh] = spg_type
            for v in spg_type.relations.values():
                v_zh, _, object_type = self.relation_info_en[type_en][v.name]
                re_list.append(v.name_zh)
                self.relations_mapper[v.name_zh] = [spg_type, object_type]
        self.schema_list = self.multischema_split_by_num(self.split_num, re_list)


class OneKE_KGPrompt(OneKEPrompt):
    template_zh: str = "你是一个图谱实体知识结构化专家。根据输入实体类型(entity type)的schema描述，从文本中抽取出相应的实体实例和其属性信息，不存在的属性不输出, 属性存在多值就返回列表，并输出为可解析的json格式。"
    template_en: str = "You are an expert in structured knowledge systems for graph entities. Based on the schema description of the input entity type, you extract the corresponding entity instances and their attribute information from the text. Attributes that do not exist should not be output. If an attribute has multiple values, a list should be returned. The results should be output in a parsable JSON format."

    def __init__(
        self,
        entity_types: List[SPGTypeName],
        language: str = "zh",
        with_description: bool = False,
        split_num: int = 4,
    ):
        super().__init__(
            types_list=entity_types,
            language=language,
            with_description=with_description,
            split_num=split_num,
        )

    def parse_response(self, response: str) -> List[SPGRecord]:
        if isinstance(response, list) and len(response) > 0:
            response = response[0]
        try:
            re_obj = json.loads(response)
        except json.decoder.JSONDecodeError:
            print("DeepKE_KGPrompt response JSONDecodeError error.")
            return []
        if type(re_obj) != dict:
            print("DeepKE_KGPrompt response type error.")
            return []

        spg_records = []
        for type_zh, type_value in re_obj.items():
            if type_zh not in self.spg_type_schema_info_zh:
                print(f"Unrecognized entity_type: {type_zh}")
                continue
            type_en, _ = self.spg_type_schema_info_zh[type_zh]
            if type_value and isinstance(type_value, dict):
                for name, attrs in type_value.items():
                    spg_record = SPGRecord(type_en).upsert_properties(
                        {"id": name, "name": name}
                    )
                    for attr_zh, attr_value in attrs.items():
                        if isinstance(attr_value, list):
                            attr_value = ",".join(attr_value)
                        if attr_zh in self.property_info_zh[type_zh]:
                            attr_en, _, object_type = self.property_info_zh[type_zh][
                                attr_zh
                            ]
                            spg_record.upsert_property(attr_en, attr_value)
                        elif attr_zh in self.relation_info_zh[type_zh]:
                            attr_en, _, object_type = self.relation_info_zh[type_zh][
                                attr_zh
                            ]
                            spg_record.upsert_relation(attr_en, object_type, attr_value)
                        else:
                            print(f"Unrecognized attribute: {attr_zh}")
                            continue
                        if object_type == "Integer":
                            matches = re.findall(r"\d+", attr_value)
                            if matches:
                                spg_record.upsert_property(attr_en, matches[0])
                        elif object_type == "Float":
                            matches = re.findall(r"\d+(?:\.\d+)?", attr_value)
                            if matches:
                                spg_record.upsert_property(attr_en, matches[0])
                    spg_records.append(spg_record)
        return spg_records

    def _render(self):
        spo_list = []
        for spg_type in self.spg_types:
            attributes = []
            attributes.extend(
                [
                    v.name_zh
                    for k, v in spg_type.properties.items()
                    if k not in ["id", "description", "stdId"]
                ]
            )
            attributes.extend(
                [
                    v.name_zh
                    for k, v in spg_type.relations.items()
                    if v not in attributes and k not in ["isA"]
                ]
            )
            entity_type = spg_type.name_zh
            spo_list.append({"entity_type": entity_type, "attributes": attributes})

        self.schema_list = self.multischema_split_by_num(self.split_num, spo_list)


class OneKE_EEPrompt(OneKEPrompt):
    template_zh: str = "你是专门进行事件提取的专家。请从input中抽取出符合schema定义的事件，不存在的事件返回空列表，不存在的论元返回NAN，如果论元存在多值请返回列表。请按照JSON字符串的格式回答。"
    template_en: str = "You are an expert in event extraction. Please extract events from the input that conform to the schema definition. Return an empty list for events that do not exist, and return NAN for arguments that do not exist. If an argument has multiple values, please return a list. Respond in the format of a JSON string."

    def __init__(
        self,
        event_types: List[SPGTypeName],
        language: str = "zh",
        with_description: bool = False,
        split_num: int = 4,
    ):
        super().__init__(
            types_list=event_types,
            language=language,
            with_description=with_description,
            split_num=split_num,
        )

    def parse_response(self, response: str) -> List[SPGRecord]:
        if isinstance(response, list) and len(response) > 0:
            response = response[0]
        try:
            ee_obj = json.loads(response)
        except json.decoder.JSONDecodeError:
            print("DeepKE_EEPrompt response JSONDecodeError error.")
            return []
        if type(ee_obj) != dict:
            print("DeepKE_EEPrompt response type error.")
            return []

        spg_records = []
        for type_zh, type_values in ee_obj.items():
            if type_zh not in self.spg_type_schema_info_zh:
                print(f"Unrecognized event_type: {type_zh}")
                continue
            type_en, _ = self.spg_type_schema_info_zh[type_zh]
            if type_values and isinstance(type_values, list):
                for type_value in type_values:
                    spg_record = SPGRecord(type_en).upsert_property("name", type_zh)
                    arguments = type_value.get("arguments")
                    if arguments and isinstance(arguments, dict):
                        for attr_zh, attr_value in arguments.items():
                            if isinstance(attr_value, list):
                                attr_value = ",".join(attr_value)
                            if attr_zh in self.property_info_zh[type_zh]:
                                attr_en, _, object_type = self.property_info_zh[
                                    type_zh
                                ][attr_zh]
                                spg_record.upsert_property(attr_en, attr_value)
                            elif attr_zh in self.relation_info_zh[type_zh]:
                                attr_en, _, object_type = self.relation_info_zh[
                                    type_zh
                                ][attr_zh]
                                spg_record.upsert_relation(
                                    attr_en, object_type, attr_value
                                )
                            else:
                                print(f"Unrecognized attribute: {attr_zh}")
                                continue
                            if object_type == "Integer":
                                matches = re.findall(r"\d+", attr_value)
                                if matches:
                                    spg_record.upsert_property(attr_en, matches[0])
                            elif object_type == "Float":
                                matches = re.findall(r"\d+(?:\.\d+)?", attr_value)
                                if matches:
                                    spg_record.upsert_property(attr_en, matches[0])
                    spg_records.append(spg_record)
        return spg_records

    def _render(self):
        event_list = []
        for spg_type in self.spg_types:
            arguments = []
            arguments.extend(
                [
                    v.name_zh
                    for k, v in spg_type.properties.items()
                    if k not in ["id", "name", "description"]
                ]
            )
            arguments.extend(
                [
                    v.name_zh
                    for k, v in spg_type.relations.items()
                    if v.name_zh not in arguments
                ]
            )
            event_type = spg_type.name_zh
            event_list.append(
                {"event_type": event_type, "trigger": True, "arguments": arguments}
            )
        self.schema_list = self.multischema_split_by_num(self.split_num, event_list)
