import os
from textwrap import dedent

from huggingface_hub import InferenceClient
from dotenv import load_dotenv

load_dotenv()
api_token = os.environ["HF_TOKEN"]
endpoint_url = os.environ["HF_API_ENDPOINT"]

# Initialize the client with your endpoint URL and token.
client = InferenceClient(
    model=endpoint_url,
    token=api_token
)

class Phi3InstructGraph:
    def __init__(self, model = "EmergentMethods/Phi-3-mini-4k-instruct-graph"):
        """
        Initialize the Phi3InstructGraph with a specified model.
        """

        self.model_path = model

    def _generate(self, messages):
        """
        Generate a response from the model based on the provided messages.
        """

        # Use the chat_completion method
        response = client.chat_completion(
            messages=messages,
            max_tokens=1024,
        )

        # Access the generated text
        generated_text = response.choices[0].message.content
        return generated_text

    def _get_messages(self, text):
        """
        Construct the message list for the chat model.
        """
        
        context = dedent("""\n
                    A chat between a curious user and an artificial intelligence Assistant. The Assistant is an expert at identifying entities and relationships in text. The Assistant responds in JSON output only.

                    The User provides text in the format:

                    -------Text begin-------
                    <User provided text>
                    -------Text end-------

                    The Assistant follows the following steps before replying to the User:

                    1. **identify the most important entities** The Assistant identifies the most important entities in the text. These entities are listed in the JSON output under the key "nodes", they follow the structure of a list of dictionaries where each dict is:

                    "nodes":[{"id": <entity N>, "type": <type>, "detailed_type": <detailed type>}, ...]

                    where "type": <type> is a broad categorization of the entity. "detailed type": <detailed_type>  is a very descriptive categorization of the entity.

                    2. **determine relationships** The Assistant uses the text between -------Text begin------- and -------Text end------- to determine the relationships between the entities identified in the "nodes" list defined above. These relationships are called "edges" and they follow the structure of:

                    "edges":[{"from": <entity 1>, "to": <entity 2>, "label": <relationship>}, ...]

                    The <entity N> must correspond to the "id" of an entity in the "nodes" list.

                    The Assistant never repeats the same node twice. The Assistant never repeats the same edge twice.
                    The Assistant responds to the User in JSON only, according to the following JSON schema:

                    {"type":"object","properties":{"nodes":{"type":"array","items":{"type":"object","properties":{"id":{"type":"string"},"type":{"type":"string"},"detailed_type":{"type":"string"}},"required":["id","type","detailed_type"],"additionalProperties":false}},"edges":{"type":"array","items":{"type":"object","properties":{"from":{"type":"string"},"to":{"type":"string"},"label":{"type":"string"}},"required":["from","to","label"],"additionalProperties":false}}},"required":["nodes","edges"],"additionalProperties":false}
                        """)
        
        user_message = dedent(f"""\n
                    -------Text begin-------
                    {text}
                    -------Text end-------
                    """)
        
        # if self.model_path == "EmergentMethods/Phi-3-medium-128k-instruct-graph":
        messages = [ 
            {
                "role": "system", 
                "content": context
            }, 
            {
                "role": "user", 
                "content": user_message
            }
        ]
        # else:
        #     # TODO: update for other models

        return messages
    
    def extract(self, text):
        """
        Extract knowledge graph from text
        """

        messages = self._get_messages(text)
        generated_text = self._generate(messages)

        return generated_text