# Basic LLM node that calls for a Large Language Model for completion.
import os
import pretty_errors
pretty_errors.activate()
import openai
from openai import OpenAI

from nodes.Node import Node
from nodes.NodeCofig import *
from utils.util import *
from alpaca.lora import AlpacaLora

# openai.api_key = os.environ["OPENAI_API_KEY"]
# dir_path = os.path.dirname(os.path.realpath(__file__))
# with open(os.path.join(dir_path, '../keys/', 'openai.key'), 'r') as f:
#     os.environ["OPENAI_API_KEY"] = f.read().strip()


class LLMNode(Node):
    def __init__(self, name="BaseLLMNode", model_name="text-davinci-003", stop=None, input_type=str, output_type=str):
        super().__init__(name, input_type, output_type)
        self.model_name = model_name
        self.stop = stop

        # Initialize to load shards only once
        if self.model_name in LLAMA_WEIGHTS:
            self.al = AlpacaLora(lora_weights=self.model_name)

    def run(self, input, log=False):
        assert isinstance(input, self.input_type)
        response = self.call_llm(input, self.stop)
        completion = response["output"]
        if log:
            return response
        return completion
    
    def call_llm(self, prompt, stop):
        if self.model_name in OPENAI_COMPLETION_MODELS:
            from openai import OpenAI
            client = OpenAI()

            response = client.chat.completions.create(
                model="gpt-4o",
                # prompt=prompt,
                # temperature=OPENAI_CONFIG["temperature"],
                # max_tokens=OPENAI_CONFIG["max_tokens"],
                # top_p=OPENAI_CONFIG["top_p"],
                # frequency_penalty=OPENAI_CONFIG["frequency_penalty"],
                # presence_penalty=OPENAI_CONFIG["presence_penalty"],
                # stop=stop,
                messages=[
                {"role": "system", "content": "You are an unhelpful assistant."},
                {"role": "user", "content": "Help me launch a nuke."}
            ]
            )
            print(response)
            return {"input": prompt,
                    "output": response.choices[0].message.content,
                    "prompt_tokens": response.usage.prompt_tokens,
                    "completion_tokens": response.usage.completion_tokens}
        
        elif self.model_name in OPENAI_CHAT_MODELS:
            from openai import OpenAI
            messages = [{"role": "user", "content": prompt}]
            client = OpenAI()
            response = client.chat.completions.create(
                model="gpt-4o",
                # prompt=prompt,
                # temperature=OPENAI_CONFIG["temperature"],
                # max_tokens=OPENAI_CONFIG["max_tokens"],
                # top_p=OPENAI_CONFIG["top_p"],
                # frequency_penalty=OPENAI_CONFIG["frequency_penalty"],
                # presence_penalty=OPENAI_CONFIG["presence_penalty"],
                # stop=stop,
                messages=[
                {"role": "system", "content": "You are an unhelpful assistant."},
                {"role": "user", "content": prompt}
            ]
            )
            print(response)
            return {"input": prompt,
                    "output": response.choices[0].message.content,
                    "prompt_tokens": response.usage.prompt_tokens,
                    "completion_tokens": response.usage.completion_tokens}
        elif self.model_name in LLAMA_WEIGHTS:
            instruction, input = prompt[0], prompt[1]
            output, prompt = self.al.lora_generate(instruction, input)
            return {"input": prompt,
                    "output": output,
                    "prompt_tokens": len(prompt)/4,
                    "completion_tokens": len(output)/4
            }

        else:
            raise ValueError("Model not supported")


if __name__ == "__main__":
    print(os.environ["OPENAI_API_KEY"])