from src.Llms.Modelos.Base_model import ModeloABC
from src.Llms.Memory.Base_MemoryConv import MemoryConvABC
from src.Llms.MCPs.McpClient_Registry import ClientRegistry
from datetime import datetime
from typing import Optional, List, Union, AsyncGenerator
import re
import json

from entrypoint.init_db import db_credencial
from src.Logger.logger_db import LoggerDB, logger
LoggerDB(db_credencial, "logger_agentes", created_by="sistema")




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",
        mcp: Optional[ClientRegistry] = None,
        output_schema: Optional[dict] = None,
    ):
        self.modelo = modelo
        self.memoria = memoria
        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
        self.mcp = mcp  # <-- Aquí guardamos el registry



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




    async def generar_system_prompt(self) -> str:

        info = f"""Eres un agente de texto y te llamas {self.nombre} 

        ### Descripción:
        {self.descripcion}

        ### Rol:
        {self.rol}

        ### Objetivos:
        {chr(10).join(f"- {o}" for o in self.objetivos)}

        ### System Prompt:
        {self.system_prompt}

        Siempre estructura tus respuestas con claridad, y termina con <END> cuando hayas completado la tarea principal del usuario.
        """.strip()

        return info





    async def construir_prompt_usuario(self, prompt_usuario: str) -> str:
        bloques = []

        if self.mcp:
            tools_str = await self._obtener_herramientas_disponibles_str()
            bloques.append(f"### Herramientas disponibles (MCP):\n{tools_str}")
            bloques.append("""### Instrucciones para actuar con herramientas MCP:
            Eres un agente conversacional con acceso a herramientas MCP. Cuando el usuario te haga una solicitud, sigue este proceso paso a paso:
            ---
            🧠 **Piensa**:  
            Reflexiona en voz alta. Explica claramente qué crees que se necesita hacer y por qué.
            🎯 **Decide**:  
            Elige si puedes resolverlo tú solo, si necesitas más información del usuario, o si una herramienta MCP sería útil.
            ⚙️ **Actúa**:  
            Si decides usar una herramienta, **escribe el bloque MCP justo después**, sin ningún texto extra después del bloque.
            ---
            ### Formato MCP:
                           
            ```mcp
            {
            "server": "tools",
            "tool": "get_current_user",
            "input": {}
            }
            ```

            ---
                           
            ### ❗ REGLAS IMPORTANTES:

            - **Puedes pensar y decidir con texto normal**, pero:
            - El **bloque MCP debe ser lo último** que aparece en tu mensaje.
            - **NO escribas nada después del bloque MCP.**
            - Solo usa `<END>` cuando:
            - hayas terminado completamente la tarea del usuario,
            - e interpretado la salida de las herramientas que usaste.
            - Puedes hacer múltiples pasos si es necesario: usar una herramienta, esperar su salida, analizarla, usar otra, etc.
            - Si decides no usar herramientas, simplemente responde como lo harías normalmente.
            - Si no estás seguro de algo, **pide aclaraciones al usuario** antes de actuar.

            ---

            ✅ Correcto:
            ```mcp
            {
            "server": "tools",
            "tool": "generate_uuid",
            "input": {}
            }
            ````
            🔵 Siempre usa ` ```mcp ` (con triple backtick y la palabra `mcp`) antes del JSON. No escribas nada después del bloque.  
            ````
            ---

            ### ✅ Ejemplo correcto:

            Necesito generar un identificador único para el usuario.  
            Para eso usaré la herramienta `generate_uuid` disponible.  
            
            ```mcp
            {
            "server": "tools",
            "tool": "generate_uuid",
            "input": {}
            }

            """)

        if self.memoria:
            historial = self.memoria.cargar_historial_chat()
            if historial:
                memoria_str = "\n".join(
                    [f"{msg['role']}: {msg['content']}" for msg in historial]
                )
                bloques.append(f"### Memoria del chat:\n{memoria_str}")

        if self.output_schema:
            schema_str = str(self.output_schema)
            bloques.append(f"### Salida esperada:\n{schema_str}")

        bloques.append(f"### Prompt del usuario:\n{prompt_usuario}")

        return "\n\n".join(bloques)
    




### Conseguir las herramientas disponibles

    async def _obtener_herramientas_disponibles_str(self) -> str:
        logger.info("Inicio de obtención de herramientas disponibles")

        if not self.mcp:
            logger.warning("No se ha definido el cliente MCP.")
            return "No se han definido herramientas disponibles."

        try:
            resultado = await self.mcp.listar_tools_por_cliente()
            tools_por_cliente = resultado.get("tools", {})
            errores = resultado.get("errores", {})

            logger.debug(f"Tools obtenidas: {list(tools_por_cliente.keys())}")
            logger.debug(f"Errores detectados: {list(errores.keys())}")

            herramientas = []

            for name, tools in tools_por_cliente.items():
                if not tools:
                    logger.info(f"Servidor {name} no tiene herramientas disponibles.")
                    continue

                herramientas.append(f"\n🔌 Server: {name}")
                for tool in tools:
                    props = tool.inputSchema.get("properties", {})
                    parametros = "\n     ".join(f"- {k} ({v.get('type', '?')})" for k, v in props.items())
                    herramientas.append(f"""Nombre: {tool.name}
            Descripción: {tool.description}
            Parámetros:
                {parametros}
            """)
                    logger.debug(f"Herramienta agregada: {tool.name} del servidor {name}")

            if errores:
                herramientas.append("\n⚠️ Los siguientes servidores no están disponibles:")
                for name, error in errores.items():
                    herramientas.append(f"- {name}: {error}")
                    logger.warning(f"Servidor con error: {name} -> {error}")

            logger.info("Finalización de obtención de herramientas exitosamente.")
            return "\n".join(herramientas) or "No hay herramientas disponibles actualmente."

        except Exception as e:
            logger.error(f"Error inesperado al obtener herramientas: {str(e)}", exc_info=True)
            return "Se produjo un error al obtener las herramientas disponibles."




### Formatear prompt para agentes

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



### Ejecutar codigo MCP

    async def ejecutar_bloque_mcp(self, respuesta: str) -> Optional[str]:
        logger.info("Iniciando ejecución de bloque MCP.")

        patron = r"```mcp\s*(\{.*?\})\s*```"
        match = re.search(patron, respuesta, re.DOTALL)

        if not match:
            patron_incorrecto = r"```[\s]*\{.*?\}[\s]*```"
            if re.search(patron_incorrecto, respuesta, re.DOTALL):
                logger.warning("Bloque detectado sin especificador `mcp`.")
                return "Advertencia: Usaste un bloque de herramienta MCP pero olvidaste indicar el lenguaje `mcp`. Corrige el bloque a: ```mcp { ... } ```"
            logger.info("No se encontró ningún bloque MCP en la respuesta.")
            return None

        try:
            bloque_json_str = match.group(1)
            logger.debug(f"Bloque MCP detectado: {bloque_json_str}")

            bloque = json.loads(bloque_json_str)

            server_name = bloque["server"]
            tool_name = bloque["tool"]
            input_args = bloque.get("input", {})

            logger.info(f"Bloque MCP válido. Servidor: {server_name}, Herramienta: {tool_name}")
            logger.debug(f"Parámetros de entrada: {input_args}")

        except Exception as e:
            logger.error(f"Error al interpretar el bloque MCP: {e}", exc_info=True)
            return f"Error al interpretar el bloque MCP: {e}"

        try:
            cliente_mcp = self.mcp.get(server_name)
        except KeyError:
            logger.warning(f"No se encontró el cliente MCP para el servidor '{server_name}'.")
            return f"No se encontró el cliente MCP para el servidor '{server_name}'"

        try:
            logger.info(f"Ejecutando herramienta '{tool_name}' en servidor '{server_name}' con argumentos: {json.dumps(input_args, ensure_ascii=False)}")

            async with cliente_mcp:
                resultado = await cliente_mcp.call_tool(tool_name, input_args)
                logger.info(f"Ejecución completada exitosamente. Resultado: {resultado}")
                return str(resultado)

        except Exception as e:
            logger.error(f"Error al ejecutar herramienta '{tool_name}' en servidor '{server_name}': {e}", exc_info=True)
            return f"Error al ejecutar herramienta '{tool_name}' en servidor '{server_name}': {e}"



### Ejecutar VARIOS bloques MCP

    async def ejecutar_multiples_bloques_mcp(self, respuesta: str) -> Optional[List[str]]:
        logger.info("Buscando múltiples bloques MCP en la respuesta.")

        patron = r"```mcp\s*(\{.*?\})\s*```"
        matches = re.finditer(patron, respuesta, re.DOTALL)

        resultados = []
        hubo_bloques = False

        for match in matches:
            hubo_bloques = True
            bloque_json_str = match.group(1)
            try:
                bloque = json.loads(bloque_json_str)
                server_name = bloque["server"]
                tool_name = bloque["tool"]
                input_args = bloque.get("input", {})

                logger.info(f"Ejecutando bloque MCP: servidor={server_name}, herramienta={tool_name}")

                try:
                    cliente_mcp = self.mcp.get(server_name)
                except KeyError:
                    msg = f"No se encontró el cliente MCP para el servidor '{server_name}'"
                    logger.warning(msg)
                    resultados.append(msg)
                    continue

                async with cliente_mcp:
                    resultado = await cliente_mcp.call_tool(tool_name, input_args)
                    resultado_str = f"[{server_name}.{tool_name}] → {resultado}"
                    resultados.append(resultado_str)

            except Exception as e:
                error_msg = f"Error al procesar bloque MCP: {str(e)}"
                logger.error(error_msg, exc_info=True)
                resultados.append(error_msg)

        if not hubo_bloques:
            logger.info("No se encontró ningún bloque MCP en la respuesta.")
            return None

        return resultados




###----------- Funcion para interactuar

    async def interactuar(self, prompt: str, stream: bool = False) -> Union[str, AsyncGenerator[str, None]]:
        mensaje_usuario = await self.construir_prompt_usuario(prompt)
        contexto = [{"role": "user", "content": mensaje_usuario}]
        prompt_final = self._formatear_prompt(contexto)

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

        return respuesta






###----------- Funcion para interactuar en bucle

    async def interactuar_en_bucle(self, prompt: str, stream: bool = False) -> Union[List[str], AsyncGenerator[str, None]]:
        respuestas = [] if not stream else None
        respuesta_anterior = ""
        resultado_mcp_anterior = None  # <-- Guarda último resultado del MCP
        iteration = 0
        prompt_original = prompt.strip()

        async def generador():
            nonlocal iteration, respuesta_anterior, resultado_mcp_anterior

            while self.max_iterations == 0 or iteration < self.max_iterations:
                instruccion_fin = (
                    "\n\nIMPORTANTE: Cuando hayas respondido completamente a la pregunta original del usuario y no requieras más pasos, "
                    "escribe <END> para indicar que has terminado."
                )

                if iteration == 0:
                    prompt_actual = prompt_original + instruccion_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"
                        f"{instruccion_fin}"
                    )

                if resultado_mcp_anterior:
                    prompt_actual += (
                        "\n\nEsta fue la salida de la herramienta que usaste:\n"
                        f"{resultado_mcp_anterior}\n\n"
                        "Úsala para seguir resolviendo el problema o tomar una nueva decisión."
                    )

                mensaje_usuario = await self.construir_prompt_usuario(prompt_actual)
                contexto = [{"role": "user", "content": mensaje_usuario}]
                prompt_final = self._formatear_prompt(contexto)

                respuesta = await self.modelo.responder(
                    prompt=prompt_final,
                    system_prompt=await self.generar_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

                # Revisar y ejecutar bloque MCP si existe
                resultado_mcp_anterior = None
                if "```mcp" in respuesta_anterior:
                    resultados_mcp = await self.ejecutar_multiples_bloques_mcp(respuesta_anterior)
                    if resultados_mcp:
                        resultado_mcp_anterior = "\n".join(resultados_mcp)

                        if stream:
                            yield "\n" + resultado_mcp_anterior
                        else:
                            respuestas.append(resultado_mcp_anterior)

                # Guardar historial si hay memoria
                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 "<end>" in respuesta_anterior.lower() and "```mcp" not in respuesta_anterior.lower():
                    break

                iteration += 1

        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]