# app/lib/embedding_providers/openai_provider.py
from typing import List
from openai import OpenAI
from .base import BaseProvider


class OpenAIProvider(BaseProvider):
    """
    OpenAI provider for embeddings and completions.
    Supports text-embedding-3-small, text-embedding-3-large, and GPT models.
    """

    EMBEDDING_DIMENSIONS = {
        "text-embedding-3-small": 1536,
        "text-embedding-3-large": 3072,
        "text-embedding-ada-002": 1536,
    }

    def __init__(
        self,
        api_key: str,
        embedding_model: str = "text-embedding-3-small",
        completion_model: str = "gpt-4o-mini"
    ):
        """
        Initialize OpenAI provider.

        Args:
            api_key: OpenAI API key
            embedding_model: Model to use for embeddings
            completion_model: Model to use for completions
        """
        if not api_key:
            raise ValueError("OpenAI API key is required")

        self.client = OpenAI(api_key=api_key)
        self.embedding_model = embedding_model
        self.completion_model = completion_model

        if embedding_model not in self.EMBEDDING_DIMENSIONS:
            raise ValueError(
                f"Unsupported embedding model: {embedding_model}. "
                f"Supported models: {list(self.EMBEDDING_DIMENSIONS.keys())}"
            )

    @property
    def dimension(self) -> int:
        return self.EMBEDDING_DIMENSIONS[self.embedding_model]

    @property
    def provider_name(self) -> str:
        return "openai"

    @property
    def model_name(self) -> str:
        return self.embedding_model

    def get_embeddings(self, texts: List[str]) -> List[List[float]]:
        """
        Generate embeddings using OpenAI API.

        Args:
            texts: List of text strings to embed

        Returns:
            List of embedding vectors

        Raises:
            ValueError: If texts is empty
            RuntimeError: If API call fails
        """
        if not texts:
            raise ValueError("Text list cannot be empty")

        try:
            response = self.client.embeddings.create(
                model=self.embedding_model,
                input=texts
            )
            return [item.embedding for item in response.data]
        except Exception as e:
            raise RuntimeError(f"OpenAI embedding API call failed: {str(e)}")

    def get_completion(self, prompt: str, context: str = "") -> str:
        """
        Generate a completion using OpenAI Chat API.

        Args:
            prompt: The user's question
            context: Optional context from retrieved documents

        Returns:
            The generated response

        Raises:
            ValueError: If prompt is empty
            RuntimeError: If API call fails
        """
        if not prompt:
            raise ValueError("Prompt cannot be empty")

        try:
            messages = []

            # Add system message with context if provided
            if context:
                messages.append({
                    "role": "system",
                    "content": f"Answer the user's question based on the following context:\n\n{context}"
                })

            # Add user message
            messages.append({
                "role": "user",
                "content": prompt
            })

            response = self.client.chat.completions.create(
                model=self.completion_model,
                messages=messages,
                max_tokens=2000,
                temperature=0.7
            )

            return response.choices[0].message.content

        except Exception as e:
            raise RuntimeError(f"OpenAI completion API call failed: {str(e)}")
