"""
Diagram Agent API — FastAPI route for AI diagram generation.
Uses DiagramGraph 2-Agent (Planner → Tool → Responder).
History stored in Redis, auto-expires after 30 minutes.
"""

import json
import logging
from pydantic import BaseModel
from fastapi import APIRouter
from fastapi.responses import JSONResponse

from langchain_core.messages import AIMessage, HumanMessage

from agent.diagram_agent.diagram_graph import get_diagram_agent
from common.cache import redis_cache

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/diagram", tags=["Diagram Agent"])

HISTORY_KEY_PREFIX = "diagram:hist:"
HISTORY_TTL = 1800  # 30 min


class DiagramChatRequest(BaseModel):
    query: str
    session_id: str | None = None


# ─── Helpers: serialize/deserialize LangChain messages ↔ Redis ───

def _serialize_messages(messages: list) -> str:
    data = []
    for msg in messages:
        content = msg.content if isinstance(msg.content, str) else str(msg.content)
        if isinstance(msg, HumanMessage):
            data.append({"role": "human", "content": content})
        elif isinstance(msg, AIMessage):
            data.append({"role": "ai", "content": content})
    return json.dumps(data, ensure_ascii=False)


def _deserialize_messages(raw: str) -> list:
    try:
        data = json.loads(raw)
    except Exception:
        return []
    messages = []
    for item in data:
        if item["role"] == "human":
            messages.append(HumanMessage(content=item["content"]))
        elif item["role"] == "ai":
            messages.append(AIMessage(content=item["content"]))
    return messages


async def _load_history(session_id: str) -> list:
    try:
        client = redis_cache.get_client()
        if not client:
            return []
        raw = await client.get(f"{HISTORY_KEY_PREFIX}{session_id}")
        if raw:
            return _deserialize_messages(raw)
    except Exception as e:
        logger.warning("Redis load history error: %s", e)
    return []


async def _save_history(session_id: str, messages: list):
    try:
        client = redis_cache.get_client()
        if not client:
            return
        trimmed = messages[-30:]
        raw = _serialize_messages(trimmed)
        await client.setex(f"{HISTORY_KEY_PREFIX}{session_id}", HISTORY_TTL, raw)
    except Exception as e:
        logger.warning("Redis save history error: %s", e)


async def _clear_history(session_id: str):
    try:
        client = redis_cache.get_client()
        if client:
            await client.delete(f"{HISTORY_KEY_PREFIX}{session_id}")
    except Exception as e:
        logger.warning("Redis clear history error: %s", e)


# ─── Endpoints ───

@router.post("/chat", summary="Chat with Diagram Agent")
async def diagram_chat(req: DiagramChatRequest):
    """Send a message to the Diagram Agent. Returns AI response + Mermaid diagram data."""
    query = req.query.strip()
    if not query:
        return JSONResponse(status_code=400, content={"status": "error", "message": "Query trống"})

    session_id = req.session_id

    try:
        agent = get_diagram_agent()

        history = []
        if session_id:
            history = await _load_history(session_id)

        result = await agent.chat(query, history=history if history else None)

        if session_id:
            history.append(HumanMessage(content=query))
            if result.get("response"):
                history.append(AIMessage(content=result["response"]))
            await _save_history(session_id, history)

        return {
            "status": "success",
            "response": result["response"],
            "elapsed_ms": result["elapsed_ms"],
            "agent_path": result["agent_path"],
            "tool_calls": result["tool_calls"],
            "pipeline": result.get("pipeline", []),
            "diagram": result.get("diagram"),
            "session_id": session_id,
            "history_count": len(history),
        }
    except Exception as e:
        logger.error(f"❌ Diagram Agent error: {e}", exc_info=True)
        return JSONResponse(status_code=500, content={"status": "error", "message": str(e)})


@router.post("/clear", summary="Clear diagram chat history")
async def diagram_clear(req: DiagramChatRequest):
    if req.session_id:
        await _clear_history(req.session_id)
    return {"status": "success", "message": "History cleared"}
