# backend/src/agent/app.py
# #region AgentChat.GradioApp [C:4] [TYPE Module] [SEMANTICS agent-chat,gradio,app]
# @DEFGROUP AgentChat Gradio ChatInterface wrapping LangGraph agent. Streaming via submit(), HITL via interrupt().
# @PRE JWT_SECRET env var set. Shared with FastAPI for stateless validation.
# @POST Agent streams tokens via Gradio yield; audit logged via LoggingMiddleware.
# @SIDE_EFFECT Calls LLM, invokes tools via FastAPI REST, writes checkpoints to PostgreSQL.
# @RELATION DEPENDS_ON -> [AgentChat.Document.Parser]
# @RELATION DEPENDS_ON -> [AgentChat.Document.Parser]
# @RELATION DEPENDS_ON -> [AgentChat.Document.Parser]
# @RELATION DEPENDS_ON -> [AgentChat.Document.Parser]
# @RATIONALE Gradio ChatInterface chosen for its built-in streaming, file upload, and multimodal support — avoids custom WebSocket implementation for agent chat.
# @REJECTED Custom React chat frontend rejected — Gradio provides free authentication, session management, and mobile-responsive UI out of the box.

from collections.abc import AsyncGenerator
import json
import os
import uuid

import gradio as gr
import httpx
import jwt
from langchain_core.exceptions import OutputParserException
from langchain_core.messages import HumanMessage
from langgraph.types import Command

from src.agent.context import set_user_jwt
from src.agent.document_parser import parse_upload
from src.agent.langgraph_setup import create_agent
from src.agent.middleware import log_tool_event
from src.agent.tools import get_all_tools
from src.core.cot_logger import log

JWT_SECRET = os.getenv("JWT_SECRET", "super-secret-key")
MAX_FILE_SIZE_BYTES = 10 * 1024 * 1024  # 10 MB

# In-memory per-user lock (keyed by user_id)
_user_locks: dict[str, bool] = {}
# In-memory service JWT cache
_service_jwt_cache: dict[str, str] = {}  # {token: expiry_timestamp}


# #region AgentChat.GradioApp.Handler [C:4] [TYPE Function] [SEMANTICS agent-chat,handler,streaming]
# @ingroup AgentChat
# @BRIEF Core streaming handler — runs LangGraph agent, yields ChatMessage tokens with structured metadata.
# @PRE JWT valid, user authenticated.
# @POST Tokens streamed via yield; HITL interrupts yield confirm_required metadata.
# @SIDE_EFFECT Calls LLM, invokes tools, writes checkpoints.
# @RATIONALE Async generator pattern chosen for Gradio ChatInterface compatibility — Gradio iterates
#            the generator and sends yielded JSON strings as event data to the frontend.
# @REJECTED Returning a single response (non-streaming) was rejected — violates FR-003 (streaming mandate).

async def agent_handler(  # noqa: C901 — intentionally complex C4 orchestration
    message,
    history: list,  # noqa: ARG001 — Gradio ChatInterface requires this parameter
    request: gr.Request,
    conversation_id: str | None = None,
    action: str | None = None,
) -> AsyncGenerator[str]:
    """Handle incoming chat message. Streams tokens with structured metadata.

    Args:
        message: str or dict (when multimodal) — user message.
        history: list of ChatMessage — Gradio's built-in history (ignored — loaded from DB).
        request: gr.Request — may contain Authorization header with user JWT.
        conversation_id: str — via additional_inputs (thread_id for checkpointer).
        action: str — "confirm" | "deny" for HITL resume, None for normal messages.
    """
    # ── Auth: extract user JWT if available —─
    # Gradio runs behind Vite proxy which already handles auth.
    # @gradio/client does not forward Authorization headers,
    # so we don't enforce JWT here. Tool calls use SERVICE_JWT (see tools.py).
    # The JWT is only used for user-scoped features (per-user lock, conversation context).
    auth_header = request.headers.get("authorization", "")
    user_jwt_str = ""
    if auth_header.startswith("Bearer "):
        try:
            token = auth_header.split(" ")[1]
            jwt.decode(token, JWT_SECRET, algorithms=["HS256"])
            user_jwt_str = token
        except jwt.InvalidTokenError:
            pass  # Ignore invalid JWTs — fall back to default context

    # Store in ContextVar for @tool functions
    set_user_jwt(user_jwt_str)

    # ── Per-user lock (prevent concurrent sends per user) ──
    user_id = _extract_user_id(user_jwt_str) if user_jwt_str else f"anon_{conversation_id or 'default'}"
    if _user_locks.get(user_id, False):
        yield json.dumps({"metadata": {"type": "error", "code": "CONCURRENT_SEND"}})
        return
    _user_locks[user_id] = True

    try:
        # ── Handle file upload ──
        text = message.get("text", "") if isinstance(message, dict) else str(message)
        files = message.get("files", []) if isinstance(message, dict) else []

        if files:
            # File size validation
            file_path = files[0] if isinstance(files[0], str) else getattr(files[0], "name", None)
            if file_path and os.path.exists(file_path):
                file_size = os.path.getsize(file_path)
                if file_size > MAX_FILE_SIZE_BYTES:
                    yield json.dumps({
                        "content": f"❌ File exceeds 10MB limit ({file_size / 1024 / 1024:.1f} MB)",
                        "metadata": {"type": "error", "code": "FILE_TOO_LARGE", "detail": "Max file size is 10 MB"},
                    })
                    return
            parsed = parse_upload(files[0])
            text = f"{text}\n\n--- Uploaded file content ---\n{parsed}"

        # ── HITL resume path ──
        if action in ("confirm", "deny"):
            async for chunk in _handle_resume(conversation_id, action):
                yield chunk
            return

        # ── Normal send path ──
        conv_id = conversation_id or str(uuid.uuid4())
        agent = create_agent(get_all_tools())

        # Try up to 2 times: catch OutputParserException and retry with stricter prompt
        max_attempts = 2
        for attempt in range(max_attempts):
            try:
                async for event in agent.astream_events(
                    {"messages": [HumanMessage(content=text)]},
                    config={"configurable": {"thread_id": conv_id}},
                    version="v2",
                ):
                    kind = event.get("event")

                    # Audit logging for tool events
                    if kind in ("on_tool_start", "on_tool_end", "on_tool_error"):
                        await log_tool_event(event, conv_id)

                    if kind == "on_chat_model_stream":
                        chunk = event["data"]["chunk"]
                        if hasattr(chunk, "content") and chunk.content:
                            yield json.dumps({
                                "content": chunk.content,
                                "metadata": {"type": "stream_token", "token": chunk.content},
                            })

                    elif kind == "on_tool_start":
                        tool_name = event["name"]
                        yield json.dumps({
                            "content": f"🛠️ {tool_name}",
                            "metadata": {"type": "tool_start", "tool": tool_name, "input": event["data"].get("input", {})},
                        })

                    elif kind == "on_tool_end":
                        tool_name = event["name"]
                        output = event["data"].get("output", "")
                        yield json.dumps({
                            "content": f"✅ {tool_name}",
                            "metadata": {"type": "tool_end", "tool": tool_name, "output": {"result": str(output)[:500]}},
                        })

                    elif kind == "on_tool_error":
                        tool_name = event["name"]
                        err = str(event["data"].get("error", "Unknown"))
                        yield json.dumps({
                            "content": f"❌ {tool_name} — {err}",
                            "metadata": {"type": "tool_error", "tool": tool_name, "error": err},
                        })

                    elif kind == "on_chain_end" and "interrupt" in event:
                        yield json.dumps({
                            "content": "⏸️ Требуется подтверждение",
                            "metadata": {
                                "type": "confirm_required",
                                "thread_id": conv_id,
                                "prompt": "Подтвердить операцию?",
                            },
                        })
                return  # Stream ends — user confirms via second submit()
                # Success — break out of retry loop
                break

            except OutputParserException as e:
                if attempt < max_attempts - 1:
                    # Retry with stricter prompt
                    text = "Respond with valid JSON only. Previous response was malformed.\n\n" + text
                    continue
                # Final failure — yield error event
                yield json.dumps({
                    "content": "❌ Ошибка обработки ответа LLM. Пожалуйста, уточните запрос.",
                    "metadata": {"type": "error", "code": "LLM_MALFORMED_OUTPUT", "detail": str(e)},
                })

        # ── Save conversation to DB via FastAPI REST ──
        await _save_conversation(conv_id, text)

    finally:
        _user_locks[user_id] = False
# #endregion AgentChat.GradioApp.Handler


async def _handle_resume(conversation_id: str, action: str) -> AsyncGenerator[str]:
    """Resume from HITL checkpoint."""
    agent = create_agent(get_all_tools())
    if action == "confirm":
        agent.invoke(
            Command(resume={"action": "confirm"}),
            config={"configurable": {"thread_id": conversation_id}},
        )
        yield json.dumps({
            "content": "▶️ Операция подтверждена",
            "metadata": {"type": "confirm_resolved", "result": "confirmed"},
        })
    elif action == "deny":
        agent.invoke(
            Command(resume={"action": "deny"}),
            config={"configurable": {"thread_id": conversation_id}},
        )
        yield json.dumps({
            "content": "⏹️ Операция отменена",
            "metadata": {"type": "confirm_resolved", "result": "denied"},
        })


def _extract_user_id(jwt_str: str) -> str:
    try:
        payload = jwt.decode(jwt_str, JWT_SECRET, algorithms=["HS256"])
        return payload.get("sub", payload.get("user_id", "unknown"))
    except Exception:
        return "unknown"


# ── Conversation persistence ──────────────────────────────────────
SAVE_API_URL = os.getenv("FASTAPI_URL", "http://localhost:8000") + "/api/agent/conversations/save"


async def _save_conversation(conv_id: str, user_text: str) -> None:
    """Save conversation to DB via FastAPI REST.

    Called after streaming completes. Creates or updates AgentConversation.
    Uses SERVICE_JWT for auth. Failures are logged but not propagated.
    """
    try:
        service_token = os.getenv("SERVICE_JWT", "")
        headers = {"Content-Type": "application/json"}
        if service_token:
            headers["Authorization"] = f"Bearer {service_token}"

        async with httpx.AsyncClient(timeout=10) as client:
            await client.post(
                SAVE_API_URL,
                json={
                    "conversation_id": conv_id,
                    "title": user_text.strip()[:100] or "Agent conversation",
                    "user_id": "0a82894e-d144-474b-aa61-81be2643d569",
                },
                headers=headers,
            )
    except Exception as e:
        log("AgentChat.GradioApp", "EXPLORE", "Failed to save conversation",
            {"conv_id": conv_id}, error=str(e))


# ── Gradio interface ──
def create_chat_interface():
    """Create the Gradio ChatInterface."""
    return gr.ChatInterface(
        fn=agent_handler,
        type="messages",
        multimodal=True,
        additional_inputs=[
            gr.Textbox(label="conversation_id", visible=False),
            gr.Textbox(label="action", visible=False),
        ],
        examples=[
            ["Покажи дашборды", None, None],
            ["Статус системы", None, None],
            ["Запусти миграцию", None, None],
        ],
    )


# ── Healthcheck ──
async def health():
    """Healthcheck endpoint for Docker."""
    return {"status": "ok", "uptime": os.times().elapsed if hasattr(os.times(), "elapsed") else 0}


if __name__ == "__main__":
    demo = create_chat_interface()
    demo.launch(
        server_name=os.getenv("GRADIO_SERVER_NAME", "0.0.0.0"),
        server_port=int(os.getenv("GRADIO_SERVER_PORT", "7860")),
    )
# #endregion AgentChat.GradioApp