import openai
import os

class Enhancer:
    def __init__(self):
        # Load the API key from an environment variable to keep it secret
        self.api_key = os.getenv('KEY')
        if not self.api_key:
            raise ValueError("API Key not found. Set KEY environment variable.")
        openai.api_key = self.api_key

    def enhance_query(self, schema, question, generated_query):
        """
        Enhances the generated query using AI techniques.
        
        Parameters:
        schema (str): The schema structure to work with.
        question (str): The question that the user asked.
        generated_query (str): The initial query generated by some model.
        
        Returns:
        str: A refined query based on the input.
        """

        # Prompt formulation for OpenAI's model
        prompt = (
            f"Schema Description:\n{schema}\n\n"
            f"User's Question:\n{question}\n\n"
            f"Initial Generated Query:\n{generated_query}\n\n"
            "Please refine the generated query to ensure that it accurately represents the schema and correctly answers the user's question."
        )

        try:
            response = openai.ChatCompletion.create(
                model="gpt-4",
                messages=[
                    {"role": "system", "content": "You are a helpful assistant specialized in SQL query refinement."},
                    {"role": "user", "content": prompt},
                ],
                max_tokens=200,
                temperature=0.3,
            )
            refined_query = response.choices[0].message['content'].strip()
            return refined_query

        except openai.error.OpenAIError as e:
            return f"Error with OpenAI API: {str(e)}"
