from src.Llms.Modelos.Base_model import ModeloABC
from src.Security.GenerarIDs import GeneradorIDUnico
from typing import AsyncGenerator, Union
from src.ConexionApis.Ollama_cliente import OllamaCliente  # Asegúrate de importar correctamente
import asyncio

class ModeloOllama(ModeloABC):
    def __init__(
        self,
        cliente: OllamaCliente,
        model: str = "llama3",
        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
    ):
        if not isinstance(cliente, OllamaCliente):
            raise TypeError("El parámetro 'cliente' debe ser una instancia de OllamaCliente")


        self.id = id if id else GeneradorIDUnico("MOOL").generar()
        super().__init__(
            model=model,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            frecuencia_penalizacion=frecuencia_penalizacion,
            num_tokens_maximos=num_tokens_maximos
        )
        self.cliente = cliente

    async def responder(
        self,
        prompt: str,
        system_prompt: str = "",
        stream: bool = False,
        **kwargs
    ) -> Union[str, AsyncGenerator[str, None]]:
        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
