# 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
from datetime import datetime
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
            # Save conversation after HITL resume
            await _save_conversation(conversation_id or str(uuid.uuid4()), "HITL resume", user_id)
            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
        try:
            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": "Подтвердить операцию?",
                                },
                            })
                    break  # Stream ends — break out to save conversation

                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)},
                    })

        except Exception:
            # Non-LLM-recoverable error (e.g. APIConnectionError).
            # Save conversation (at least user message) before re-raising.
            await _save_conversation(conv_id, text, user_id)
            raise

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

    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, user_id: str = "admin") -> None:
    """Save conversation to DB via FastAPI REST.

    Called after streaming completes. Creates or updates AgentConversation
    and persists messages. 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}"

        payload = {
            "conversation_id": conv_id,
            "title": user_text.strip()[:100] or "Agent conversation",
            "user_id": user_id,
            "messages": [
                {
                    "id": str(uuid.uuid4()),
                    "conversation_id": conv_id,
                    "role": "user",
                    "text": user_text.strip(),
                    "created_at": datetime.utcnow().isoformat(),
                }
            ],
        }

        async with httpx.AsyncClient(timeout=10) as client:
            await client.post(SAVE_API_URL, json=payload, 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
