from openai import AsyncOpenAI, OpenAI
from pydantic import BaseModel
from typing import Type
from scraipe.extended.llm_analyzers.llm_analyzer_base import LlmAnalyzerBase

class OpenAiAnalyzer(LlmAnalyzerBase):
    """An analyzer that integrates the OpenAI API.
    
    This analyzer makes asynchronous calls to the OpenAI API to obtain language model completions.
    """

    def __init__(self,
        api_key: str,
        instruction: str,
        organization: str = None,
        pydantic_schema: Type[BaseModel] = None,
        model: str = "gpt-4o-mini",
        max_content_size: int = 10000,
        max_workers: int = 3):
        """Initializes the OpenAiAnalyzer instance.
        
        Args:
            api_key (str): The API key for accessing the OpenAI API.
            instruction (str): The instruction for guiding the language model's responses.
            organization (str, optional): The organization identifier for the OpenAI API. Defaults to None.
            pydantic_schema (Type[BaseModel], optional): A Pydantic model for validating the API response structure. Defaults to None.
            model (str, optional): The model identifier to be used for generating completions. Defaults to "gpt-4o-mini".
            max_content_size (int, optional): The maximum length of the content to analyze. Defaults to 10000.
            max_workers (int, optional): The maximum number of workers for concurrent analysis. Defaults to 3.
        """
        super().__init__(
            instruction=instruction, pydantic_schema=pydantic_schema,
            max_content_size=max_content_size, max_workers=max_workers)
        self.api_key = api_key
        self.organization = organization
        self.client = AsyncOpenAI(api_key=api_key, organization=organization)
        self.model = model
        
        # Connect to OpenAi synchronously to ensure the API key and model are valid
        self.validate()
        
    def validate(self) -> None:
        test_client = OpenAI(api_key=self.api_key, organization=self.organization)
        model = test_client.models.retrieve(model=self.model)
        assert model is not None, f"Model {model} not found in OpenAI API. Please check your API key and model name."
    
    async def query_llm(self, content: str, instruction: str) -> str:
        """Asynchronously queries the OpenAI API using the provided content and instruction.
        
        The method constructs a message list with a system instruction and user-provided content,
        sends the request to the OpenAI API, and extracts the resulting message content.
        
        Args:
            content (str): The user-provided content for analysis.
            instruction (str): The system-level instruction for the language model.
            
        Returns:
            str: The content of the response message generated by the language model.
        """
        messages = [
            {"role": "system", "content": instruction},
            {"role": "user", "content": content}
        ]
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            response_format={ "type": "json_object" }
        )
        response_content: str = response.choices[0].message.content
        return response_content