import os
from openai import OpenAI
from promptsy.prompt import Prompt
from promptsy.prompt_manager import PromptManager

class PromptEnhancer:
    def __init__(self, api_key=None, model_name="gpt-4o-mini"):
        """
        Initializes the PromptEnhancer with the API key and model name.

        :param api_key: OpenAI API key. If None, tries to retrieve from the environment variable.
        :param model_name: The name of the model to be used (default: "gpt-4o-mini").
        """
        if api_key is None:
            api_key = os.getenv("OPENAI_API_KEY")
            if api_key is None:
                raise ValueError("OpenAI API key not provided and OPENAI_API_KEY environment variable not set.")
        
        self.client = OpenAI(api_key=api_key)
        self.model_name = model_name
        self.prompt_manager = PromptManager()
        self.enhanced_prompts_directory = 'enhanced_prompts'
        os.makedirs(self.enhanced_prompts_directory, exist_ok=True)

    def enhance_prompt(self, prompt: Prompt):
        """
        Enhances a given prompt by generating an improved version of the prompt.

        :param prompt: A Prompt object containing the original prompt template.
        :return: A new Prompt object with the enhanced template.
        """
        # Extract the prompt template from the Prompt object
        prompt_template = prompt.template
       
        # Combine the enhancement prompt with the original prompt
        intent_prompt = self.prompt_manager.load("promptsy.get_users_intent")
        
        summary_intent = self._call_llm(intent_prompt.format(input_prompt=prompt_template))
        
        # Get the enhancement prompt from Promptsy
        enhancement_prompt = self.prompt_manager.load("promptsy.enhancement_prompt")

        # Call the LLM to generate the enhanced prompt
        enhanced_prompt_template = self._call_llm(enhancement_prompt.format(original_prompt=prompt_template, summary_of_intention=summary_intent))
       
        # Create a new Prompt object with the enhanced prompt template
        enhanced_prompt = Prompt(
            name=prompt.name,
            description=prompt.description,
            template=enhanced_prompt_template
        )

        self.save_enhanced_prompt(enhanced_prompt)
        return enhanced_prompt

    def _call_llm(self, prompt):
        """
        Calls the language model (LLM) with the provided prompt and returns the response.

        :param prompt: The prompt to be sent to the LLM.
        :return: The response generated by the LLM.
        """
        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=[{"role": "user", "content": prompt}],
            stream=False,
        )
        return response.choices[0].message.content.strip()

    def save_enhanced_prompt(self, prompt: Prompt):
        """
        Saves the enhanced prompt using the PromptManager.

        :param prompt: A Prompt object containing the enhanced prompt template.
        """
        self.prompt_manager.save(prompt, self.enhanced_prompts_directory+"."+prompt.name)  # Uses the save method of PromptManager
        print(f"Enhanced prompt saved using PromptManager: {prompt.name}")
