import json
import re
import requests
import pandas as pd
from abc import ABC, abstractmethod
from typing import Dict, List
# from .utils.openai import openai_call
# from .utils.logger import get_logger
# from .utils.prompts import REVIEWER, EDITOR, CREATE_NODE
from utils.openai import openai_call
from utils.logger import get_logger
from utils.prompts import REVIEWER, EDITOR, CREATE_NODE

class GraphBuilder(ABC):
    logger = get_logger("build graph")
    
    @abstractmethod
    def preprocess(self, *args):
        '''prepare data from graph extraction'''
        ...
    
    @abstractmethod
    def extractor(self):
        '''extract initial entities and relations from the content'''
        ...
        
    @abstractmethod
    def build(self):
        ...
        
    def get_chat_completion(self, messages) -> str | None:
        for i in range(3):
            try:
                reply = openai_call(messages=messages)
                assert reply, "failed to get reply."
                return reply
            except requests.exceptions.RequestException as e:
                self.logger.error(f"Connection error occurred when request for chat completion: {e}")
            except:
                self.logger.info(f"{i + 1}/3 - try to get reply for message: {messages}")
            
    def get_json_chat_completion(self, messages) -> Dict | None:
        def load_json_result(content):
            # tools
            def extract_json_res(content: str) -> Dict:
                json_marker_pattern = r"```json([\s\S]*?)```"
                match = re.search(json_marker_pattern, content, re.DOTALL)
                json_content = match.group(1)
                dict_content = json.loads(json_content)
                return dict_content
            def llm_json_formatter(content: str) -> str:
                prompt = f'''Please check if the text conforms to JSON format. If it does not, output the correct JSON format result or extract the part in JSON format; if it does, return the original text.
                    Please return a valid JSON result, without any extra explanations or symbols.
                    TEXT: {content}'''
                res = openai_call(
                    messages=[{"role": "user", "content": prompt.format(content)}],
                    llm_model_name='gpt-4o-mini'
                )
                return res
            
            # main workflow
            try:
                json_res = json.loads(content)
                return json_res
            except:
                pass
            try:
                json_res = extract_json_res(content=content)
                return json_res
            except Exception as e:
                self.logger.debug(f"failed to directly load the content: {content}. ERROR: {e}. try to use LLM to format it.")
            try:
                new_content = llm_json_formatter(content=content)
                json_res = extract_json_res(new_content)
                return json_res
            except Exception as e:
                self.logger.debug(f"the content: {content}, improved by the LLM: {new_content}, still not conform to the JSON format. ERROR: {e}")
        for _ in range(3):
            try:
                _reply = self.get_chat_completion(messages=messages)
                reply = load_json_result(_reply)
                assert reply, "failed to get json reply."
                return reply
            except:
                pass
        self.logger.error(f"failed to get result for message: {messages}")
        # TODO: maybe raise error
    
    def checker(self, content=None):
        '''
        Check whether the entities correspond to the relationships and whether there are discrete entity nodes or duplicate entity nodes.
        - if there is any entity mentioned in some relation but not exist in the entity list, create a corresponding entity node.
        - if there is any discrete entity node, delete it.
        - if there are duplicate entity nodes, only keep the first one.
        '''
        def create_node(entity, context):
            mess = [
                {"role": "user", "content": CREATE_NODE.format(entity=entity, content=context)}
            ]
            node = self.get_json_chat_completion(messages=mess)
            if not node:
                self.logger.warning(f"failed to create detailed entity node for newly emerged entity: {entity}, a basic node will be inserted instead.")
                return {
                    "entity name": entity,
                    "entity type": None,
                    "timestamp": None,
                    "description": None
                }
            else:
                return node
        
        self.logger.info("Check the entity list and relation list...")
        # check for duplicate
        self.entity = self.entity.drop_duplicates(subset='entity name', keep='first')
        
        entity_name = self.entity['entity name']
        rel_entity1 = self.relation['entity1']
        rel_entity2 = self.relation['entity2']
        rel_entity = pd.concat([rel_entity1, rel_entity2], ignore_index=True).drop_duplicates().reset_index(drop=True)
        new_entity = rel_entity[~rel_entity.isin(entity_name)]
        discrete_entity = entity_name[~entity_name.isin(rel_entity)]
        # check for newly merged entity
        if new_entity.empty:
            self.logger.info("entity list matches the relation list.")
        else:
            self.logger.info(f"Found mismatched entities: {', '.join(new_entity.values)}")
            for e in new_entity.values:
                temp_node = create_node(e, content)
                self.entity.loc[len(self.entity)] = temp_node
            self.logger.info("successfully completed the entity list.")
        # check for discrete entity
        if discrete_entity.empty:
            self.logger.info("no discrete entity found.")
        else:
            self.logger.info(f"found discrete entities: {', '.join(discrete_entity.values)}")
            self.entity.drop(discrete_entity.index, inplace=True)
            self.logger.info("successfully remove the discrete entities.")
        
            
        
    def imporve(self, content):
        self.logger.info("Trying to improve the quality of the graph...")
        result = f"ENTITY: {json.dumps(self.entity.to_dict(orient='records'))}\nRELATION: {json.dumps(self.relation.to_dict(orient='records'))}"
        reviewer_mes = [
            {
                "role": "user", 
                "content": REVIEWER.format(title=self.title, content=content, result=result)
            }
        ]
        review = self.get_chat_completion(reviewer_mes)
        if review:
            if not 'PASS' in review:
                self.logger.info(f"Optimization suggestions from the reviewer:\n{review}")
                writer_mes = [
                    {
                        "role": "user", 
                        "content": EDITOR.format(review=review, history=result)
                    }
                ]
                modified_res = self.get_json_chat_completion(writer_mes)
                if modified_res:
                    try:
                        new_entity, new_relation = list(modified_res.values())
                        # new_entity = modified_res['entity']
                        # new_relation = modified_res['relation']
                        self.entity = pd.DataFrame(new_entity)
                        self.relation = pd.DataFrame(new_relation)
                        self.logger.info(f"Successfully completed the graph optimization process. The updated version:\n{json.dumps(modified_res, indent=4)}")
                        return True
                    except Exception as e:
                        self.logger.warning(f"The editor failed to generate a valid optimization result due to {e}. Raw content:\n{json.dumps(modified_res, indent=4)}")
                else:
                    self.logger.warning("The editor failed to generate a valid optimization result.")
                        
            else:
                self.logger.info("The reviewer believes that the current figure is already sufficiently refined and suggests skipping the optimization process.")
        else:
            self.logger.warning("The reviewer was unable to provide constructive feedback.")
    