"""{{ ui_title | default('ABI') }} — Chainlit Chat Interface.

A thin client over an ABI agent's ``/stream`` (SSE) endpoint. It opens a
framework-managed session so multi-turn conversations stay stable (a clarifying
question and its answer land in the same context), and streams status/result
updates into the Chainlit UI. Plan-confirmation replies (approve/reject/modify)
are rendered as real action buttons when the agent emits
``meta.action_type == "plan_confirmation"``.

Generated by ABI-Core scaffolding. Run with: chainlit run app.py
"""

import ast
import json
import warnings

import chainlit as cl
import httpx

# Suppress httpcore async generator cleanup warnings (cosmetic, non-blocking)
warnings.filterwarnings("ignore", message=".*async generator ignored GeneratorExit.*")

from config import AGENT_URL

# Sentinel queries the orchestrator's classify_query step recognizes as a
# plan-confirmation reply (see abi_agents/orchestrator/agent/steps.py).
_PLAN_CONFIRM_SENTINELS = {
    "approve": "__plan_confirm_approve__",
    "reject": "__plan_confirm_reject__",
    "modify": "__plan_confirm_modify__",
}


def _parse_message(msg):
    """Parse the SSE 'message' payload into a dict, safely.

    The stream is JSON, but the inner ``message`` can arrive as a JSON object or
    as a Python-dict string (``{'response_type': 'text', ...}``) depending on how
    the agent serialized it. Try JSON first, then ``ast.literal_eval`` (safe —
    never executes code, unlike ``eval``).
    """
    if isinstance(msg, dict):
        return msg
    if not isinstance(msg, str):
        return {"content": str(msg)}
    try:
        return json.loads(msg)
    except (json.JSONDecodeError, ValueError):
        pass
    try:
        parsed = ast.literal_eval(msg)
        return parsed if isinstance(parsed, dict) else {"content": str(msg)}
    except (ValueError, SyntaxError):
        return {"content": msg}


async def _ensure_session() -> str | None:
    """Start a session once and cache its token for this chat.

    Returns the bearer token, or ``None`` if the agent doesn't expose
    ``/session/start`` (older agents) — in which case we fall back to a
    tokenless request.
    """
    token = cl.user_session.get("abi_session_token")
    if token:
        return token
    try:
        async with httpx.AsyncClient(timeout=30) as client:
            resp = await client.post(f"{AGENT_URL}/session/start")
            if resp.status_code == 200:
                token = resp.json().get("session_token")
                cl.user_session.set("abi_session_token", token)
                return token
    except Exception:
        pass
    return None


@cl.on_chat_start
async def on_start():
    await _ensure_session()
    await cl.Message(content="🐝 **{{ ui_title | default('ABI') }}** ready. What do you want to build?").send()


async def _stream_to_agent(query: str) -> None:
    """Send `query` to the agent's /stream endpoint and render the response.

    Shared by on_message (free-text turns) and the plan-confirmation action
    callback (button clicks) — both are just a query string to the agent, the
    server doesn't distinguish "action" from "text" (see web_interface.py's
    /stream contract).
    """
    response_msg = cl.Message(content="")
    await response_msg.send()

    status_step = cl.Step(name="Working", type="run")
    await status_step.send()

    token = await _ensure_session()
    headers = {"Authorization": f"Bearer {token}"} if token else {}

    final_content = ""
    actions_sent = False

    try:
        async with httpx.AsyncClient(timeout=600) as client:
            try:
                async with client.stream(
                    "POST",
                    f"{AGENT_URL}/stream",
                    json={"query": query},
                    headers=headers,
                ) as response:
                    async for line in response.aiter_lines():
                        if not line.startswith("data: "):
                            continue
                        raw = line[6:].strip()
                        if not raw or raw == "{}":
                            continue
                        try:
                            data = json.loads(raw)
                        except json.JSONDecodeError:
                            continue

                        parsed = _parse_message(data.get("message", ""))
                        if not parsed:
                            continue

                        response_type = parsed.get("response_type", "")
                        content = parsed.get("content", "")

                        if response_type == "status" and content:
                            status_step.output = content
                            await status_step.update()
                        elif response_type == "text" and content:
                            final_content = content
                        elif response_type == "data" and content:
                            # Final result payload (AgentResponse.result). Render
                            # something human-readable: pull "result" if present,
                            # otherwise stringify the payload.
                            if isinstance(content, dict):
                                final_content = str(content.get("result", content))
                            else:
                                final_content = str(content)
                        elif response_type == "input_required" and content:
                            meta = parsed.get("meta", {}) or {}
                            if meta.get("action_type") == "plan_confirmation":
                                actions = [
                                    cl.Action(name="plan_confirm", payload={"choice": "approve"}, label="✅ Aprobar"),
                                    cl.Action(name="plan_confirm", payload={"choice": "reject"}, label="❌ Rechazar"),
                                    cl.Action(name="plan_confirm", payload={"choice": "modify"}, label="✏️ Modificar"),
                                ]
                                await cl.Message(content=content, actions=actions).send()
                                actions_sent = True
                            else:
                                # The agent needs clarification (or plain
                                # input); surface as text. The next message
                                # continues the same session (same token).
                                final_content = content
                        elif response_type == "error" and content:
                            final_content = f"❌ {content}"
            except (RuntimeError, GeneratorExit):
                pass  # Stream closed by client — ignore

    except httpx.ConnectError:
        final_content = "❌ Cannot connect to the agent at {{ ui_title | default('ABI') }}. Is it running?"
    except Exception as e:
        final_content = f"❌ Error: {str(e)}"

    status_step.output = "Done"
    await status_step.update()

    if actions_sent:
        try:
            await response_msg.remove()
        except Exception:
            pass
    else:
        response_msg.content = final_content or "No response received."
        await response_msg.update()


@cl.on_message
async def on_message(message: cl.Message):
    """Send the user message to the agent and stream the response."""
    await _stream_to_agent(message.content)


@cl.action_callback("plan_confirm")
async def on_plan_confirm(action: cl.Action):
    """Handle a plan-confirmation button click (approve/reject/modify)."""
    choice = action.payload.get("choice")
    await _stream_to_agent(_PLAN_CONFIRM_SENTINELS.get(choice, _PLAN_CONFIRM_SENTINELS["reject"]))
