from typing import List, Dict, Optional
from .base_llm import BaseLLM
from openai import OpenAI
import os
import base64

class OpenAILlm(BaseLLM):
    def __init__(self, model: str = "gpt-4o", api_key: Optional[str] = None):
        """
        Initialize the OpenAI LLM class.

        Args:
            model (str): The name of the model (e.g., gpt-4o, gpt-4-vision).
            api_key (Optional[str]): The OpenAI API key. If not provided, it fetches from the environment.
        """
        self.model = model
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
        if not self.api_key:
            raise ValueError("OpenAI API key is required. Set OPENAI_API_KEY environment variable or pass it explicitly.")
        self.client = OpenAI(api_key=self.api_key)

    def generate(self, prompt: str, context: Optional[List[Dict]] = None, memory: Optional[List[Dict]] = None) -> str:
        """
        Generate text using OpenAI's ChatCompletion API.

        Args:
            prompt (str): The user prompt.
            context (Optional[List[Dict]]): Context to include in the conversation.
            memory (Optional[List[Dict]]): Memory from previous interactions.

        Returns:
            str: The generated response from the model.
        """
        try:
            # Prepare messages for the OpenAI API
            messages = []
            if memory:
                messages.extend(memory)
            if context:
                messages.append({"role": "system", "content": "Context: " + str(context)})
            messages.append({"role": "user", "content": prompt})

            # Call the ChatCompletion endpoint
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
            )

            # Extract and return the response:
            return response.choices[0].message.content
        except Exception as e:
            raise ValueError(f"Error while generating response with OpenAI: {e}")

    @property
    def supports_vision(self) -> bool:
        """
        Check if the model supports vision tasks.
        """
        # List of GPT models that support vision
        vision_models =[
            "gpt-4o", "gpt-4o mini", "o1", "o1 mini"
        ]
        return self.model in vision_models

    def generate_from_image_url(self, prompt: str, image_url: str, **kwargs) -> str:
        """
        Process an image URL with OpenAI's vision-capable models, using instructions as the prompt.

        Args:
            image_url (str): The URL of the image.
            instructions (str): Instructions provided as the prompt for image analysis.
            kwargs: Additional parameters for the OpenAI API.

        Returns:
            str: The response generated by the vision-capable model.
        """
        if not self.supports_vision:
            raise ValueError(f"Model '{self.model}' does not support vision tasks.")

        try:
            # Use instructions as the prompt in the API call
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": prompt},  # Using instructions as the prompt
                            {"type": "image_url", "image_url": {"url": image_url,},},
                        ],
                    }
                ],
                **kwargs,
            )
            return response.choices[0].message.content
        except Exception as e:
            raise ValueError(f"Error while processing image URL with OpenAI Vision model: {e}")

    def generate_from_image(self, prompt: str, image_bytes: bytes, **kwargs) -> str:
        """
        Process an image and generate a response if the model supports vision.
        """
        if not self.supports_vision:
            raise ValueError(f"Model '{self.model}' does not support vision tasks.")

        try:
            # Convert the image bytes to base64
            image_base64 = base64.b64encode(image_bytes).decode("utf-8")
            
            # Construct the message payload
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{image_base64}",
                            },
                        },
                    ],
                }
            ]

            # Call the Groq API with the base64-encoded image
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                **kwargs,
            )

            # Extract and return the response text
            return response.choices[0].message.content
        except Exception as e:
            raise ValueError(f"Error while processing image with OpenAI vision model: {e}")
