"""nogic mcp — stdio server that proxies tool calls to the Nogic backend.

Architecture:
    AI client (Cursor/Claude/etc.)  <->  stdio (JSON-RPC)  <->  nogic mcp  <->  HTTP  <->  api.nogic.dev

The CLI acts as a thin proxy:
  1. Forwards tool calls to the backend HTTP MCP endpoint
  2. Intercepts responses containing `write_to` fields
  3. Writes files locally (agent canvases, reviews) before returning the result
"""

import asyncio
import json
import os
import sys
import uuid
from pathlib import Path
from typing import Any

import httpx
import typer

from nogic.config import Config


def _log(msg: str) -> None:
    """Log to stderr (stdout is reserved for MCP protocol messages)."""
    print(f"[nogic-mcp] {msg}", file=sys.stderr, flush=True)


# ---------------------------------------------------------------------------
# Backend HTTP MCP client
# ---------------------------------------------------------------------------


class _BackendClient:
    """Lightweight JSON-RPC client that speaks MCP over Streamable HTTP."""

    def __init__(self, url: str, api_key: str):
        self.url = url
        self.headers: dict[str, str] = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "Accept": "application/json, text/event-stream",
        }
        self._http = httpx.AsyncClient(timeout=300.0)
        self._next_id = 0
        self._reinitializing = False

    def _id(self) -> int:
        self._next_id += 1
        return self._next_id

    async def _post(self, body: dict, *, _retried: bool = False) -> dict:
        """Send a JSON-RPC request and return the parsed response."""
        resp = await self._http.post(self.url, json=body, headers=self.headers)

        # On 400, the session may be stale — re-initialize once and retry
        if resp.status_code == 400 and not _retried and not self._reinitializing:
            _log(f"Got 400 (body: {resp.text[:200]}), re-initializing session…")
            self.headers.pop("mcp-session-id", None)
            self._reinitializing = True
            try:
                await self.initialize()
            finally:
                self._reinitializing = False
            return await self._post(body, _retried=True)

        resp.raise_for_status()

        # Capture MCP session ID from the first response
        sid = resp.headers.get("mcp-session-id")
        if sid:
            self.headers["mcp-session-id"] = sid

        ct = resp.headers.get("content-type", "")
        if "text/event-stream" in ct:
            # SSE fallback: extract the last JSON-RPC data frame
            last: dict = {}
            for line in resp.text.splitlines():
                if line.startswith("data:"):
                    payload = line[5:].strip()
                    if payload:
                        try:
                            last = json.loads(payload)
                        except json.JSONDecodeError:
                            pass
            if not last:
                _log(f"SSE response contained no valid data frames (status={resp.status_code})")
            return last

        return resp.json()

    # -- MCP lifecycle --

    async def initialize(self) -> dict:
        """Initialize MCP session with the backend."""
        result = await self._post(
            {
                "jsonrpc": "2.0",
                "method": "initialize",
                "params": {
                    "protocolVersion": "2025-03-26",
                    "capabilities": {},
                    "clientInfo": {"name": "nogic-cli", "version": "0.0.6"},
                },
                "id": self._id(),
            }
        )

        # Acknowledge initialization (non-critical)
        try:
            await self._http.post(
                self.url,
                json={"jsonrpc": "2.0", "method": "notifications/initialized"},
                headers=self.headers,
            )
        except httpx.HTTPError:
            pass

        return result.get("result", {})

    async def list_tools(self) -> list[dict]:
        """List all available tools from the backend."""
        data = await self._post(
            {"jsonrpc": "2.0", "method": "tools/list", "id": self._id()}
        )
        return data.get("result", {}).get("tools", [])

    async def call_tool(self, name: str, arguments: dict) -> dict:
        """Call a tool on the backend and return the raw MCP result."""
        data = await self._post(
            {
                "jsonrpc": "2.0",
                "method": "tools/call",
                "params": {"name": name, "arguments": arguments},
                "id": self._id(),
            }
        )
        if "error" in data:
            raise RuntimeError(data["error"].get("message", str(data["error"])))
        return data.get("result", {})

    async def close(self) -> None:
        await self._http.aclose()


# ---------------------------------------------------------------------------
# File-writing middleware
# ---------------------------------------------------------------------------


def _is_within(child: Path, parent: Path) -> bool:
    """Check that *child* is inside *parent*, immune to prefix collisions."""
    parent_str = str(parent.resolve()) + os.sep
    return str(child.resolve()).startswith(parent_str) or child.resolve() == parent.resolve()


def _resolve_project_data_dir(config: Config) -> Path | None:
    """Resolve the centralized project data directory (~/.nogic/projects/{project_id}/)."""
    if not config.project_id:
        return None
    try:
        uuid.UUID(config.project_id)
    except (ValueError, AttributeError):
        _log(f"Blocked invalid project_id in path resolution: {config.project_id}")
        return None
    return Path.home() / ".nogic" / "projects" / config.project_id


def _rewrite_path(write_to: str, project_data_dir: Path | None) -> Path | None:
    """Map a .nogic/-relative write_to path to the centralized project data dir."""
    if not project_data_dir:
        return None
    # Strip the .nogic/ prefix and write to ~/.nogic/projects/{pid}/{remainder}
    if write_to.startswith(".nogic/"):
        remainder = write_to[len(".nogic/"):]
        resolved = (project_data_dir / remainder).resolve()
        # Guard against path traversal (e.g. "../../etc/passwd")
        if not _is_within(resolved, project_data_dir):
            _log(f"Blocked path traversal attempt: {write_to}")
            return None
        return resolved
    return None


def _resolve_project_data_dir_from_id(project_id: str | None) -> Path | None:
    """Build the centralized data dir from a raw project_id string."""
    if not project_id:
        return None
    # Validate project_id is a UUID to prevent path traversal
    try:
        uuid.UUID(project_id)
    except (ValueError, AttributeError):
        _log(f"Blocked invalid project_id in path resolution: {project_id}")
        return None
    return Path.home() / ".nogic" / "projects" / project_id


def _handle_file_writes(
    data: dict,
    project_data_dir: Path | None,
    tool_arguments: dict[str, Any] | None = None,
) -> dict:
    """Intercept ``write_to`` instructions, write files, strip write metadata."""
    if not isinstance(data, dict) or "write_to" not in data:
        return data

    # If the startup config didn't resolve a project, try the tool arguments
    effective_data_dir = project_data_dir
    if not effective_data_dir and tool_arguments:
        effective_data_dir = _resolve_project_data_dir_from_id(
            tool_arguments.get("project_id")
        )

    # Rewrite to centralized storage (~/.nogic/projects/{pid}/...)
    write_path = _rewrite_path(data["write_to"], effective_data_dir)
    if not write_path:
        _log(f"Cannot resolve write path (no project data dir): {data['write_to']}")
        return data

    # agent_canvas uses ``spec``; finalize_review uses ``review``
    content = data.get("spec") or data.get("review")
    if content:
        write_path.parent.mkdir(parents=True, exist_ok=True)
        write_path.write_text(json.dumps(content, indent=2))
        data["file_written"] = str(write_path)
        _log(f"Wrote {write_path}")

    # finalize_review may also include review history
    if "history_write_to" in data and "history" in data:
        history_path = _rewrite_path(data["history_write_to"], effective_data_dir)
        if history_path:
            history_path.parent.mkdir(parents=True, exist_ok=True)
            history_path.write_text(json.dumps(data["history"], indent=2))
            data["history_file_written"] = str(history_path)
            _log(f"Wrote {history_path}")

    # Strip write instructions — the AI client doesn't need them
    for key in ("instruction", "write_to", "history_write_to"):
        data.pop(key, None)

    # Strip large payloads that have already been written to disk
    for key in ("spec", "review", "history"):
        data.pop(key, None)

    return data


# ---------------------------------------------------------------------------
# Stdio MCP server
# ---------------------------------------------------------------------------


async def _run_mcp(workspace_arg: str | None) -> None:
    from mcp.server import Server
    from mcp.server.stdio import stdio_server
    import mcp.types as types

    workspace_root = Path(workspace_arg or ".").resolve()
    config = Config.load(workspace_root)

    if not config.api_key:
        _log("No API key found. Run `nogic login` or set NOGIC_API_KEY.")
        sys.exit(1)

    project_data_dir = _resolve_project_data_dir(config)
    if project_data_dir:
        _log(f"Project data dir: {project_data_dir}")

    mcp_url = f"{config.api_url}/mcp"
    _log(f"Connecting to {mcp_url}")

    backend = _BackendClient(mcp_url, config.api_key)

    try:
        init_result = await backend.initialize()
    except Exception as exc:
        _log(f"Failed to connect to backend: {exc}")
        await backend.close()
        sys.exit(1)

    instructions = init_result.get("instructions", "")
    raw_tools = await backend.list_tools()
    _log(f"Discovered {len(raw_tools)} tools")

    # Prepend the resolved project_id so the client never has to look it up
    if config.project_id:
        instructions = (
            f"## Project ID (use this for ALL tool calls)\n"
            f"`project_id`: `{config.project_id}`\n\n"
            f"Do NOT read files to find the project_id — use the value above.\n\n"
            + instructions
        )

    # Convert raw tool dicts to mcp.types.Tool objects
    tool_objects = [
        types.Tool(
            name=t["name"],
            description=t.get("description", ""),
            inputSchema=t.get("inputSchema", {"type": "object", "properties": {}}),
        )
        for t in raw_tools
    ]

    # Create local MCP server
    server = Server("nogic", instructions=instructions)

    @server.list_tools()
    async def handle_list_tools() -> list[types.Tool]:
        return tool_objects

    @server.call_tool()
    async def handle_call_tool(
        name: str,
        arguments: dict[str, Any] | None,
    ) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
        args = arguments or {}

        # Auto-inject project_id when the client omits it or sends an invalid value
        if config.project_id:
            given = args.get("project_id")
            if not given:
                args = {**args, "project_id": config.project_id}
            else:
                try:
                    uuid.UUID(given)
                except (ValueError, AttributeError):
                    _log(f"Replacing invalid project_id '{given}' with {config.project_id}")
                    args = {**args, "project_id": config.project_id}

        result = await backend.call_tool(name, args)

        processed: list[
            types.TextContent | types.ImageContent | types.EmbeddedResource
        ] = []

        for block in result.get("content", []):
            if block.get("type") == "text":
                text = block["text"]
                try:
                    parsed = json.loads(text)
                    if isinstance(parsed, dict) and "write_to" in parsed:
                        parsed = _handle_file_writes(parsed, project_data_dir, args)
                        text = json.dumps(parsed)
                except (json.JSONDecodeError, TypeError, ValueError):
                    pass
                processed.append(types.TextContent(type="text", text=text))
            elif block.get("type") == "image":
                processed.append(
                    types.ImageContent(
                        type="image",
                        data=block.get("data", ""),
                        mimeType=block.get("mimeType", "image/png"),
                    )
                )
            else:
                # Unknown type — pass through as text
                processed.append(
                    types.TextContent(type="text", text=json.dumps(block))
                )

        if not processed:
            processed.append(
                types.TextContent(type="text", text=json.dumps(result))
            )

        return processed

    _log("Server ready (stdio)")

    try:
        async with stdio_server() as (read_stream, write_stream):
            await server.run(
                read_stream,
                write_stream,
                server.create_initialization_options(),
            )
    finally:
        await backend.close()


# ---------------------------------------------------------------------------
# CLI command
# ---------------------------------------------------------------------------


def mcp_command(
    workspace: str = typer.Option(
        None,
        "--workspace",
        "-w",
        help="Workspace root for file writes (defaults to current directory).",
    ),
) -> None:
    """Start an MCP stdio server that proxies to the Nogic backend."""
    # Force unbuffered stdio so JSON-RPC responses reach the MCP client
    # immediately. Without this, Python fully buffers stdout when spawned
    # as a subprocess (no TTY), causing the client to time out.
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(line_buffering=True)
    if hasattr(sys.stdin, "reconfigure"):
        sys.stdin.reconfigure(line_buffering=True)
    asyncio.run(_run_mcp(workspace))
