from src.Llms.Modelos.Base_model import ModeloABC
from src.Llms.Memory.Base_MemoryConv import MemoryConvABC

from datetime import datetime
from typing import Optional, List, Union, AsyncGenerator


class AgenteAI:
    def __init__(
        self,
        modelo: ModeloABC,
        nombre: str,
        descripcion: str,
        system_prompt: str,
        rol: str,
        objetivos: List[str],
        max_iterations: int = 1,
        memoria: Optional[MemoryConvABC] = None,
        version: str = "1.0.0",
        tools: Optional[List] = None,
        output_schema: Optional[dict] = None,
    ):
        self.modelo = modelo
        self.memoria = memoria
        self.tools = tools or []
        self.output_schema = output_schema

        self.nombre = nombre
        self.descripcion = descripcion
        self.system_prompt = system_prompt
        self.max_iterations = max_iterations
        self.rol = rol
        self.objetivos = objetivos
        self.version = version

        self.created_at = datetime.now()
        self.updated_at = self.created_at
        self.numero_interacciones = 0

    def actualizar_configuracion(self, **kwargs):
        for clave, valor in kwargs.items():
            if hasattr(self, clave):
                setattr(self, clave, valor)
        self.updated_at = datetime.now()

    @property
    def full_system_prompt(self) -> str:
        partes = [
            f"Tu nombre es: {self.nombre}",
            f"Tu descripción: {self.descripcion}",
            f"Tu Rol: {self.rol}",
            f"Tus Objetivos: {', '.join(self.objetivos)}",
            ""
        ]

        herramientas = self._obtener_descripcion_tools()
        if herramientas:
            partes.append("Estas son tus herramientas disponibles:")
            partes.extend(herramientas)
            partes.append(
            "Cuando consideres necesario, utiliza las herramientas disponibles "
            "para responder de manera más precisa o realizar tareas específicas. "
            "Indica claramente qué herramienta estás utilizando y por qué."
            )

        partes.append(self.system_prompt)

        if self.output_schema:
            partes.append("SIEMPRE formatea la respuesta final siguiendo estrictamente el siguiente esquema JSON:")
            partes.append(f"```json\n{self.output_schema}\n```")

        return "\n".join(partes)

    def _obtener_descripcion_tools(self) -> List[str]:
        descripciones = []
        if not hasattr(self, "mcp_servers"):
            return descripciones

        for server in self.mcp_servers:
            if hasattr(server, "tools") and server.tools:
                for tool in server.tools:
                    if isinstance(tool, str):
                        descripciones.append(f"- {tool}: [sin descripción]")
                    elif isinstance(tool, dict):
                        nombre = tool.get("name", "¿?")
                        descripcion = tool.get("description", "[sin descripción]")
                        descripciones.append(f"- {nombre}: {descripcion}")
                    elif hasattr(tool, "name"):
                        descripcion = getattr(tool, "description", "[sin descripción]")
                        descripciones.append(f"- {tool.name}: {descripcion}")
        return descripciones

    def _formatear_prompt(self, mensajes: List[dict]) -> str:
        return "\n".join([f"{msg['role']}: {msg['content']}" for msg in mensajes])

    async def interactuar(self, prompt: str, stream: bool = False) -> Union[str, AsyncGenerator[str, None]]:
        historial = self.memoria.cargar_historial_chat() if self.memoria else []
        contexto = historial + [{"role": "user", "content": prompt}]
        prompt_final = self._formatear_prompt(contexto)

        respuesta = await self.modelo.responder(
            prompt=prompt_final,
            system_prompt=self.full_system_prompt,
            stream=stream
        )

        if stream:
            # stream es un generador asincrónico
            async def wrapper():
                buffer_respuesta = ""
                async for token in respuesta:
                    buffer_respuesta += token
                    yield token
                if self.memoria:
                    self.memoria.guardar_turno("user", prompt)
                    self.memoria.guardar_turno("assistant", buffer_respuesta)
                self.numero_interacciones += 1
                self.updated_at = datetime.now()
            return wrapper()
        else:
            if self.memoria:
                self.memoria.guardar_turno("user", prompt)
                self.memoria.guardar_turno("assistant", respuesta)
            self.numero_interacciones += 1
            self.updated_at = datetime.now()
            return respuesta

    async def interactuar_en_bucle(self, prompt: str, stream: bool = False) -> Union[List[str], AsyncGenerator[str, None]]:
        historial = self.memoria.cargar_historial_chat() if self.memoria else []
        respuestas = [] if not stream else None
        respuesta_anterior = None
        iteration = 0
        prompt_original = prompt.strip()

        async def generador():
            nonlocal iteration, respuesta_anterior
            prompt_actual = prompt_original
            while self.max_iterations == 0 or iteration < self.max_iterations:
                if iteration == 0:
                    prompt_actual += (
                        "\n\nIMPORTANTE:\n"
                        "Si al revisar tu última respuesta y mi pregunta inicial consideras que has terminado, "
                        "di alguna de estas frases: <FIN>"
                    )
                else:
                    prompt_actual = (
                        f"Esta es la pregunta original:\n{prompt_original}\n\n"
                        f"Esto fue lo último que dijiste:\n{respuesta_anterior}\n"
                        
                        "\n\nIMPORTANTE:\n"
                        "Si al revisar tu última respuesta y mi pregunta inicial consideras que has terminado, "
                        "di alguna de estas frases: <FIN>"
                        "En caso contrario, responde a la pregunta original "
                        "y añade información relevante que no hayas mencionado antes.\n\n"
                    )

                contexto = historial + [{"role": "user", "content": prompt_actual}]
                prompt_final = self._formatear_prompt(contexto)

                respuesta = await self.modelo.responder(
                    prompt=prompt_final,
                    system_prompt=self.full_system_prompt,
                    stream=stream
                )

                if stream:
                    buffer_respuesta = ""
                    async for token in respuesta:
                        buffer_respuesta += token
                        yield token
                    respuesta_anterior = buffer_respuesta
                else:
                    respuestas.append(respuesta)
                    respuesta_anterior = respuesta

                if self.memoria:
                    self.memoria.guardar_turno("user", prompt_actual)
                    self.memoria.guardar_turno("assistant", respuesta_anterior)

                self.numero_interacciones += 1
                self.updated_at = datetime.now()

                if "<fin>" in respuesta_anterior.lower():
                    break

                iteration += 1
                prompt_actual = ""

        return generador() if stream else await generador_to_list(generador)

# Helper para consumir generador asincrónico si no es stream
async def generador_to_list(gen: AsyncGenerator[str, None]) -> List[str]:
    buffer = ""
    async for chunk in gen:
        buffer += chunk
    return [buffer]