"""Brain Studio (Terminal) - Interactive agent chat + pathway control.

Features:
- Chat with any agent (default: host)
- Switch agents with /agent <id>
- Live event stream via WebSocket
- Pathway streaming control plane (JSON-RPC)
"""

from __future__ import annotations

import asyncio
import json
import time
from urllib.parse import quote

import aiohttp
from aiohttp import WSMsgType


# ---------------------------------------------------------------------------
# Event formatters: event_type → human-readable log line
# ---------------------------------------------------------------------------

_EVENT_FORMATTERS = {
    "turn_failed": lambda o: f"  [error] {o.get('error', 'unknown')}",
    "tool_called": lambda o: f"  [tool] calling {o.get('tool_name', '')}...",
    "tool_completed": lambda o: f"  [tool] {o.get('tool_name', '')} done",
    "tool_failed": lambda o: f"  [tool] {o.get('tool_name', '')} FAILED",
    "node_failed": lambda o: f"  [node] {o.get('node_id', '')} failed: {o.get('error', '')}",
}


# ---------------------------------------------------------------------------
# Slash command handlers — each is a focused async function
# ---------------------------------------------------------------------------

_BREAK = "break"


class _StudioCtx:
    """Mutable state for an interactive studio session."""

    __slots__ = ("session", "base_url", "ws", "agent", "thread", "rpc_id", "pending")

    def __init__(
        self,
        session: aiohttp.ClientSession,
        base_url: str,
        ws: aiohttp.ClientWebSocketResponse,
        thread: str,
    ):
        self.session = session
        self.base_url = base_url
        self.ws = ws
        self.agent = "host"
        self.thread = thread
        self.rpc_id = 0
        self.pending: dict[int, asyncio.Future[dict[str, object]]] = {}


async def _cmd_quit(_ctx: _StudioCtx, _arg: str) -> str:
    return _BREAK


async def _cmd_clear(_ctx: _StudioCtx, _arg: str) -> str | None:
    print("\033[2J\033[H", end="")
    return None


async def _cmd_agent(ctx: _StudioCtx, arg: str) -> str | None:
    if arg:
        ctx.agent = arg
        print(f"  Switched to agent: {ctx.agent}")
    else:
        print(f"  Current agent: {ctx.agent}")
    return None


async def _cmd_agents(ctx: _StudioCtx, _arg: str) -> str | None:
    try:
        async with ctx.session.get(f"{ctx.base_url}/api/v1/agents") as resp:
            data = await resp.json()
            agents = data.get("agents", [])
            print(f"\n  Agents ({len(agents)}):")
            for a in agents:
                marker = " <--" if a.get("id") == ctx.agent else ""
                print(f"    {a.get('id')}: {a.get('name', '')}{marker}")
            print()
    except Exception as e:
        print(f"  error: {e}")
    return None


async def _cmd_thread(ctx: _StudioCtx, arg: str) -> str | None:
    if arg:
        ctx.thread = arg
        print(f"  Switched to thread: {ctx.thread}")
    else:
        print(f"  Current thread: {ctx.thread}")
    return None


async def _cmd_pathways(ctx: _StudioCtx, _arg: str) -> str | None:
    try:
        async with ctx.session.get(f"{ctx.base_url}/api/v1/pathways") as resp:
            data = await resp.json()
            pathways = data.get("pathways", [])
            print(f"\n  Pathways ({len(pathways)}):")
            for p in pathways:
                print(f"    {p.get('id')} ({p.get('source', '')})")
            print()
    except Exception as e:
        print(f"  error: {e}")
    return None


async def _cmd_tools(ctx: _StudioCtx, arg: str) -> str | None:
    category = arg.strip()
    try:
        async with ctx.session.get(f"{ctx.base_url}/api/v1/tools") as resp:
            data = await resp.json()
            tools = data.get("tools", [])
            if category:
                tools = [t for t in tools if t.get("name", "").startswith(f"{category}.")]
            print(f"\n  Tools ({len(tools)}):")
            for t in tools:
                print(f"    {t.get('name')}: {(t.get('description') or '')[:60]}")
            print()
    except Exception as e:
        print(f"  error: {e}")
    return None


async def _cmd_run(ctx: _StudioCtx, arg: str) -> str | None:
    pathway_id = arg.strip()
    if not pathway_id:
        print("  usage: /run <pathway_id>")
        return None
    ctx.rpc_id += 1
    rid = ctx.rpc_id
    fut: asyncio.Future[dict[str, object]] = asyncio.get_event_loop().create_future()
    ctx.pending[rid] = fut
    await ctx.ws.send_str(
        json.dumps(
            {
                "jsonrpc": "2.0",
                "id": rid,
                "method": "pathway.stream_start",
                "params": {"pathway_id": pathway_id, "inputs": {}},
            }
        )
    )
    try:
        result = await asyncio.wait_for(fut, timeout=15.0)
        print(f"  stream started: {result.get('stream_id')}")
    except Exception as e:
        print(f"  error: {e}")
    return None


async def _cmd_help(_ctx: _StudioCtx, _arg: str) -> str | None:
    print()
    print("  /agent <id>   - Switch agent")
    print("  /agents       - List agents")
    print("  /thread <id>  - Switch thread")
    print("  /pathways     - List pathways")
    print("  /tools [cat]  - List tools")
    print("  /run <id>     - Stream pathway")
    print("  /clear        - Clear screen")
    print("  /quit         - Exit")
    print()
    return None


_SLASH_COMMANDS: dict[str, Any] = {
    "/quit": _cmd_quit,
    "/exit": _cmd_quit,
    "/clear": _cmd_clear,
    "/agent": _cmd_agent,
    "/agents": _cmd_agents,
    "/thread": _cmd_thread,
    "/pathways": _cmd_pathways,
    "/tools": _cmd_tools,
    "/run": _cmd_run,
    "/help": _cmd_help,
}


def _ws_url(host: str, port: int, thread_id: str) -> str:
    return f"ws://{host}:{port}/api/v1/ws?mode=both&thread_id={quote(thread_id)}"


async def _send_agent_turn(
    session: aiohttp.ClientSession,
    base_url: str,
    agent_id: str,
    message: str,
    thread_id: str,
) -> dict[str, object]:
    """Send a message to an agent and return the response."""
    url = f"{base_url}/api/v1/agents/{agent_id}/turn"
    payload = {"message": message, "thread_id": thread_id}
    timeout = aiohttp.ClientTimeout(total=300)
    try:
        async with session.post(url, json=payload, timeout=timeout) as resp:
            if resp.status != 200:
                return {"success": False, "error": f"Server returned {resp.status}"}
            return await resp.json()
    except TimeoutError:
        return {"success": False, "error": "Request timed out (300s)"}
    except Exception as e:
        return {"success": False, "error": str(e)}


async def _run_studio(*, host: str, port: int, thread_id: str) -> int:
    """Run the terminal studio — interactive agent chat."""
    ws_url = _ws_url(host, port, thread_id)
    base_url = f"http://{host}:{port}"

    print("\n🧙 ALBUS STUDIO")
    print(f"   agent:  host")
    print(f"   thread: {thread_id}")
    print(f"   server: {base_url}")
    print()
    print("Type a message to chat with the agent.")
    print()
    print("Commands:")
    print("  /agent <id>   - Switch agent (e.g. /agent dev)")
    print("  /agents       - List available agents")
    print("  /thread <id>  - Switch thread")
    print("  /pathways     - List pathways")
    print("  /tools [cat]  - List tools (optional category filter)")
    print("  /run <id>     - Stream a pathway by id")
    print("  /clear        - Clear screen")
    print("  /quit, /exit  - Exit studio")
    print()

    async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=None)) as session:
        # Health check
        try:
            async with session.get(f"{base_url}/api/v1/health") as resp:
                if resp.status != 200:
                    print(f"  Server not healthy at {base_url}")
                    return 1
        except Exception as e:
            print(f"  Cannot reach server at {base_url}: {e}")
            print("  Start with: uv run albus server --debug")
            return 1

        print(f"  Connected to {base_url}\n")

        # Connect WebSocket for live events
        try:
            async with session.ws_connect(ws_url, heartbeat=15.0, autoping=True) as ws:
                ctx = _StudioCtx(session, base_url, ws, thread_id)

                async def _receiver() -> None:
                    """Receive WebSocket messages (events + RPC responses)."""
                    async for msg in ws:
                        if msg.type == WSMsgType.TEXT:
                            try:
                                obj = json.loads(msg.data)
                            except Exception:
                                continue

                            # JSON-RPC response
                            if (
                                isinstance(obj, dict)
                                and obj.get("jsonrpc") == "2.0"
                                and "id" in obj
                            ):
                                rid = obj["id"]
                                if rid in ctx.pending:
                                    fut = ctx.pending.pop(rid)
                                    if "error" in obj:
                                        fut.set_exception(
                                            RuntimeError(
                                                str(
                                                    obj["error"].get("message", obj["error"])
                                                )
                                            )
                                        )
                                    else:
                                        fut.set_result(obj.get("result", {}))
                                continue

                            # Live events — show activity while agent is thinking
                            if isinstance(obj, dict):
                                et = obj.get("type") or obj.get("event_type") or ""
                                formatter = _EVENT_FORMATTERS.get(et)
                                if formatter:
                                    print(formatter(obj))

                        elif msg.type in (
                            WSMsgType.CLOSED,
                            WSMsgType.CLOSE,
                            WSMsgType.ERROR,
                        ):
                            break

                recv_task = asyncio.create_task(_receiver())

                try:
                    while True:
                        # Read input
                        prompt = f"you ({ctx.agent})> "
                        try:
                            line = await asyncio.get_event_loop().run_in_executor(
                                None, lambda: input(prompt)
                            )
                        except EOFError:
                            break

                        text = line.strip()
                        if not text:
                            continue

                        # --------------------------------------------------
                        # Slash commands — dispatch table
                        # --------------------------------------------------
                        if text.startswith("/"):
                            parts = text.split(None, 1)
                            cmd_name = parts[0]
                            arg = parts[1] if len(parts) > 1 else ""
                            handler = _SLASH_COMMANDS.get(cmd_name)
                            if handler:
                                result = await handler(ctx, arg)
                                if result == _BREAK:
                                    break
                                continue
                            print(
                                f"  Unknown command: {cmd_name}. "
                                f"Type /help for commands."
                            )
                            continue

                        # --------------------------------------------------
                        # Agent chat — send message to current agent
                        # --------------------------------------------------
                        print("  thinking...", end="", flush=True)
                        t0 = time.monotonic()

                        response = await _send_agent_turn(
                            session,
                            base_url,
                            ctx.agent,
                            text,
                            ctx.thread,
                        )

                        elapsed = time.monotonic() - t0
                        # Clear the "thinking..." line
                        print("\r" + " " * 40 + "\r", end="")

                        if response.get("success"):
                            reply = response.get("response", "")
                            steps = response.get("steps_taken", 0)
                            tools_used = response.get("tool_calls", [])
                            duration = response.get("duration_ms", 0)

                            # Show tool calls summary if any
                            if tools_used:
                                tool_names = [
                                    tc.get("tool_name", "?")
                                    for tc in tools_used
                                    if isinstance(tc, dict)
                                ]
                                if tool_names:
                                    print(
                                        f"  [used: {', '.join(tool_names)}]"
                                    )

                            # Show the response
                            print(f"\n{ctx.agent}> {reply}\n")

                            # Compact stats
                            print(
                                f"  ({steps} steps, "
                                f"{duration or int(elapsed * 1000)}ms)"
                            )
                            print()
                        else:
                            error = response.get("error", "Unknown error")
                            print(f"\n  error: {error}\n")

                finally:
                    recv_task.cancel()
                    await asyncio.gather(recv_task, return_exceptions=True)

        except aiohttp.WSServerHandshakeError:
            # WebSocket not available — fall back to HTTP-only mode
            print("  WebSocket not available, using HTTP-only mode.\n")
            while True:
                prompt = f"you ({current_agent})> "
                try:
                    line = await asyncio.get_event_loop().run_in_executor(
                        None, lambda: input(prompt)
                    )
                except EOFError:
                    break

                text = line.strip()
                if not text:
                    continue
                if text in ("/quit", "/exit"):
                    break
                if text.startswith("/agent "):
                    current_agent = text[len("/agent ") :].strip() or current_agent
                    print(f"  Switched to agent: {current_agent}")
                    continue
                if text.startswith("/"):
                    print("  HTTP-only mode: only /agent, /quit supported.")
                    continue

                print("  thinking...", end="", flush=True)
                response = await _send_agent_turn(
                    session, base_url, current_agent, text, current_thread,
                )
                print("\r" + " " * 40 + "\r", end="")

                if response.get("success"):
                    print(f"\n{current_agent}> {response.get('response', '')}\n")
                else:
                    print(f"\n  error: {response.get('error', 'Unknown')}\n")

        except Exception as e:
            print(f"  WebSocket error: {e}")
            return 1

    print("\nGoodbye")
    return 0


def run_studio(*, host: str, port: int, thread_id: str) -> int:
    """Run terminal studio (blocking)."""
    try:
        return asyncio.run(_run_studio(host=host, port=port, thread_id=thread_id))
    except KeyboardInterrupt:
        print("\nInterrupted")
        return 0


__all__ = ["run_studio"]
