from domains.Llms.Modelos.Base_model import ModeloABC
from domains.ConexionApis.OpenAi_conexion import OpenAICliente
from domains.Security.GenerarIDs import GeneradorIDUnico
import asyncio
from typing import AsyncGenerator, Union

class ModeloOpenAI(ModeloABC):
    def __init__(
        self,
        cliente: OpenAICliente,
        model: str = "gpt-4o",
        id: str = None,
        temperature: float = 0.7,
        top_p: float = 1.0,
        top_k: int = None,
        frecuencia_penalizacion: float = 0.0,
        num_tokens_maximos: int = 512,
        use_legacy: bool = False
    ):
        # Generar ID con prefijo MOPA si no fue proporcionado
        self.id = id if id is not None else GeneradorIDUnico("MOPA").generar()

        # Inicializar resto del modelo base
        super().__init__(
            model=model,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            frecuencia_penalizacion=frecuencia_penalizacion,
            num_tokens_maximos=num_tokens_maximos
        )

        # Asignar cliente e indicadores adicionales
        self.cliente = cliente
        self.use_legacy = use_legacy

    async def responder(
        self,
        prompt: str,
        system_prompt: str = "",
        stream: bool = False,
        **kwargs
    ) -> Union[str, AsyncGenerator[str, None]]:
        if self.use_legacy:
            if stream:
                raise NotImplementedError("El modo legacy no soporta streaming.")
            
            loop = asyncio.get_event_loop()
            respuesta = await loop.run_in_executor(
                None,
                lambda: self.cliente.completion(
                    model=self.model,
                    prompt=prompt,
                    temperature=self.temperature,
                    top_p=self.top_p,
                    max_tokens=self.num_tokens_maximos,
                    frequency_penalty=self.frecuencia_penalizacion,
                    **kwargs
                )
            )
            return respuesta.choices[0].text.strip()

        # Construcción de mensajes estilo Chat
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})

        def sync_call():
            return self.cliente.chat_completion(
                model=self.model,
                messages=messages,
                temperature=self.temperature,
                top_p=self.top_p,
                max_tokens=self.num_tokens_maximos,
                frequency_penalty=self.frecuencia_penalizacion,
                stream=stream,
                **kwargs
            )

        loop = asyncio.get_event_loop()
        resultado = await loop.run_in_executor(None, sync_call)

        if stream:
            async def generador():
                for token in resultado:
                    yield token
            return generador()
        else:
            return resultado.choices[0].message.content