import os
from pathlib import Path
from datetime import datetime
import pytz
import redis.asyncio as redis
from fastapi import FastAPI, HTTPException, APIRouter
from pydantic import BaseModel
import logging
import json
from dotenv import load_dotenv

from langchain_openai import ChatOpenAI
from langchain.agents import AgentExecutor, create_openai_functions_agent
from langchain.memory import ConversationBufferMemory
from langchain_mcp_adapters.client import MultiServerMCPClient

# === IMPORTAR PROMPTS DE AGENTES ===
from agents.padel_agent import get_padel_prompt
from agents.restaurante_agent import get_restaurante_prompt

# ==========================
# CONFIGURACIÓN BASE
# ==========================
BASE_DIR = Path(__file__).resolve().parent
load_dotenv(BASE_DIR / ".env", override=True)

OPENAI_KEY = os.getenv("OPENAI_KEY")
MCP_URL = os.getenv("MCP_URL")
REDIS_URL = os.getenv("REDIS_URL")

if not OPENAI_KEY:
    raise ValueError("⚠️ Falta OPENAI_KEY en el archivo .env")

# Fecha y hora de Ecuador (Zona de Guayaquil)
EC_TZ = pytz.timezone("America/Guayaquil")
CURRENT_DATETIME_EC = datetime.now(EC_TZ)
CURRENT_DATE_EC = CURRENT_DATETIME_EC.strftime("%Y-%m-%d")
CURRENT_TIME_EC = CURRENT_DATETIME_EC.strftime("%H:%M")

# ==========================
# LOGGING CONFIG
# ==========================
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    force=True
)
logger = logging.getLogger(__name__)
logging.getLogger("httpx").setLevel(logging.WARNING)
logging.getLogger("langchain_mcp_adapters").setLevel(logging.WARNING)

# ==========================
# APP FASTAPI
# ==========================
app = FastAPI(
    title="Cliente MULTIAGENTE MCP",
    description="Endpoint para recibir mensajes de WhatsApp e interactuar con el Agente LangChain.",
    root_path="/multiagente-mcp"
)

router = APIRouter()
redis_client = redis.from_url(REDIS_URL, decode_responses=True)

llm = ChatOpenAI(
    openai_api_key=OPENAI_KEY,
    model="gpt-4o",
    temperature=0.2
)

# Memorias y agentes activos
CONVERSATION_MEMORIES = {}
AGENT_EXECUTORS = {}

# ==========================
# Pydantic Models
# ==========================
class UserMessage(BaseModel):
    whatsapp_id: str
    mensaje: str

# ==========================
# Funciones de Utilidad
# ==========================
def get_or_create_memory(user_id: str):
    """Crea o recupera la memoria de conversación del usuario."""
    if user_id not in CONVERSATION_MEMORIES:
        CONVERSATION_MEMORIES[user_id] = ConversationBufferMemory(
            memory_key="history",
            return_messages=True
        )
    return CONVERSATION_MEMORIES[user_id]

async def set_user_state(whatsapp_id: str, estado: int):
    """Guarda el estado del usuario (1=pádel, 2=restaurante)."""
    await redis_client.hset(f"whatsapp:{whatsapp_id}", mapping={"estado": estado})
    logger.info(f"💾 Estado de {whatsapp_id} guardado como {estado}")

async def get_user_state(whatsapp_id: str) -> int:
    """Obtiene el estado actual del usuario."""
    estado = await redis_client.hget(f"whatsapp:{whatsapp_id}", "estado")
    return int(estado) if estado else 0

# ==========================
# Selección de Agente
# ==========================
async def get_agent_by_state(user_id: str):
    """Devuelve el agente (pádel o restaurante) según el estado del usuario."""
    estado = await get_user_state(user_id)

    if estado == 1:
        agente = "padel"
        prompt = get_padel_prompt(CURRENT_DATE_EC, CURRENT_TIME_EC, "[RESERVA_CREADA_EXITOSAMENTE]")
    elif estado == 2:
        agente = "restaurante"
        prompt = get_restaurante_prompt(CURRENT_DATE_EC, CURRENT_TIME_EC, "[RESTAURANTE_RESERVA]")
    else:
        return None

    # 🔹 Conectar al MCP
    mcp_client = MultiServerMCPClient({agente: {"url": MCP_URL, "transport": "streamable_http"}})
    all_tools = await mcp_client.get_tools()
    logger.info(f"🔧 Herramientas detectadas en MCP ({agente}): {[t.name for t in all_tools]}")

    # Crear el agente LangChain
    executor = AgentExecutor(
        agent=create_openai_functions_agent(llm=llm, tools=all_tools, prompt=prompt),
        tools=all_tools,
        memory=None,
        verbose=True
    )

    # Envolver las tools con logs
    for tool in executor.tools:
        original_func = tool.func

        async def wrapper(*args, _tool=tool, **kwargs):
            logger.info(f"🛠️ Ejecutando tool: {_tool.name} | args: {args} kwargs: {kwargs}")
            result = await original_func(*args, **kwargs)
            logger.info(f"✅ Tool {_tool.name} ejecutada. Resultado: {json.dumps(result, ensure_ascii=False)}")
            return result

        tool.func = wrapper

    return executor

# ==========================
# ENDPOINT PRINCIPAL
# ==========================
@router.post("/message")
async def receive_message(msg: UserMessage):
    """Recibe mensajes de WhatsApp y redirige al agente correspondiente."""
    logger.info(f"👤 CLIENTE [{msg.whatsapp_id}]: {msg.mensaje}")

    memory = get_or_create_memory(msg.whatsapp_id)
    estado = await get_user_state(msg.whatsapp_id)

    # Opción para salir del sistema
    if msg.mensaje.strip().lower() == "salir":
        await set_user_state(msg.whatsapp_id, 0)
        return {"status": "ok", "respuesta": "Has salido del sistema. Para iniciar nuevamente, escribe cualquier mensaje."}

    # Menú inicial
    if estado == 0:
        menu_msg = (
            "¡Hola! 👋 Bienvenido al sistema MultiAgente.\n"
            "Por favor elige una opción:\n"
            "1️⃣ Reservar una cancha de pádel\n"
            "2️⃣ Reservar una mesa en el restaurante"
        )

        if msg.mensaje.strip() == "1":
            await set_user_state(msg.whatsapp_id, 1)
            estado = 1
            logger.info(f"🟢 Usuario {msg.whatsapp_id} seleccionó agente PÁDEL")
        elif msg.mensaje.strip() == "2":
            await set_user_state(msg.whatsapp_id, 2)
            estado = 2
            logger.info(f"🟢 Usuario {msg.whatsapp_id} seleccionó agente RESTAURANTE")
        else:
            logger.info(f"📋 Mostrando menú inicial a {msg.whatsapp_id}")
            return {"status": "ok", "respuesta": menu_msg}

    # Obtener agente correspondiente
    agent_executor = await get_agent_by_state(msg.whatsapp_id)
    if not agent_executor:
        return {"status": "ok", "respuesta": "No hay agente asignado para tu estado actual."}

    agent_executor.memory = memory
    user_input = f"[{msg.whatsapp_id}] {msg.mensaje}"

    try:
        logger.info(f"➡️ Enviando input al agente ({estado}): {user_input}")
        result = await agent_executor.ainvoke({"input": user_input})
        logger.info(f"🤖 AGENTE [{estado}] respondió: {result['output']}")
        return {"status": "ok", "respuesta": result["output"]}
    except Exception as e:
        logger.error(f"❌ Error procesando mensaje de {msg.whatsapp_id}: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Error procesando mensaje: {e}")

app.include_router(router)

# ==========================
# STARTUP
# ==========================
@app.on_event("startup")
async def startup_event():
    """Inicializa los agentes y sus herramientas al arrancar el servidor."""
    logger.info("🚀 Servidor arrancando...")

    agentes_config = {
        "padel": get_padel_prompt(CURRENT_DATE_EC, CURRENT_TIME_EC, "[RESERVA_CREADA_EXITOSAMENTE]"),
        "restaurante": get_restaurante_prompt(CURRENT_DATE_EC, CURRENT_TIME_EC, "[RESTAURANTE_RESERVA]")
    }

    all_tools_info = {}

    for agente_name, prompt in agentes_config.items():
        mcp_client = MultiServerMCPClient({agente_name: {"url": MCP_URL, "transport": "streamable_http"}})
        tools = await mcp_client.get_tools()

        executor = AgentExecutor(
            agent=create_openai_functions_agent(llm=llm, tools=tools, prompt=prompt),
            tools=tools,
            memory=None,
            verbose=True
        )

        for tool in executor.tools:
            original_func = tool.func

            async def wrapper(*args, _tool=tool, **kwargs):
                logger.info(f"🛠️ Ejecutando tool: {_tool.name} | args: {args} kwargs: {kwargs}")
                result = await original_func(*args, **kwargs)
                logger.info(f"✅ Tool {_tool.name} ejecutada. Resultado: {json.dumps(result, ensure_ascii=False)}")
                return result

            tool.func = wrapper

        AGENT_EXECUTORS[agente_name] = executor
        all_tools_info[agente_name] = [t.name for t in tools]

    logger.info(f"🔍 Herramientas cargadas para todos los agentes: {all_tools_info}")
    logger.info("✅ Agentes inicializados y listos para recibir mensajes.")