#!/usr/bin/env python3
"""
Demo: make an MCP request against mcplocal using an McpToken bearer.

This is the standalone counterpart to `mcpctl test mcp` — intended to show
exactly what a non-Claude client (e.g. a vLLM-driven agent) would do.

Usage:
    # Default: localhost mcplocal, sre project, token from $MCPCTL_TOKEN
    export MCPCTL_TOKEN=mcpctl_pat_...
    python3 scripts/demo-mcp-call.py

    # Custom URL/project/tool
    python3 scripts/demo-mcp-call.py \\
        --url https://mcp.ad.itaz.eu \\
        --project sre \\
        --token "$MCPCTL_TOKEN" \\
        --tool begin_session \\
        --args '{"description":"hello"}'

No third-party deps — pure stdlib. Mirrors the protocol that
src/shared/src/mcp-http/index.ts implements on the TypeScript side.
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import urllib.error
import urllib.request
from typing import Any


def _parse_sse(body: str) -> list[dict[str, Any]]:
    """Parse a text/event-stream body into a list of JSON-RPC messages."""
    out: list[dict[str, Any]] = []
    for line in body.splitlines():
        if line.startswith("data: "):
            try:
                out.append(json.loads(line[6:]))
            except json.JSONDecodeError:
                pass
    return out


class McpSession:
    def __init__(self, url: str, bearer: str | None = None, timeout: float = 30.0):
        self.url = url
        self.bearer = bearer
        self.timeout = timeout
        self.session_id: str | None = None
        self._next_id = 1

    def _headers(self) -> dict[str, str]:
        h = {
            "Content-Type": "application/json",
            "Accept": "application/json, text/event-stream",
        }
        if self.bearer:
            h["Authorization"] = f"Bearer {self.bearer}"
        if self.session_id:
            h["mcp-session-id"] = self.session_id
        return h

    def send(self, method: str, params: dict[str, Any] | None = None) -> Any:
        rid = self._next_id
        self._next_id += 1
        payload = {"jsonrpc": "2.0", "id": rid, "method": method, "params": params or {}}
        req = urllib.request.Request(
            self.url,
            data=json.dumps(payload).encode("utf-8"),
            headers=self._headers(),
            method="POST",
        )
        try:
            with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                body = resp.read().decode("utf-8")
                content_type = resp.headers.get("content-type", "")
                # First successful response carries the session id.
                if self.session_id is None:
                    sid = resp.headers.get("mcp-session-id")
                    if sid:
                        self.session_id = sid
                messages: list[dict[str, Any]] = (
                    _parse_sse(body) if "text/event-stream" in content_type else [json.loads(body)]
                )
        except urllib.error.HTTPError as e:
            err_body = e.read().decode("utf-8", errors="replace")
            raise SystemExit(f"HTTP {e.code} from {self.url}: {err_body}") from None
        except urllib.error.URLError as e:
            raise SystemExit(f"transport error reaching {self.url}: {e.reason}") from None

        # Pick the response matching our id; fall back to first message.
        matched = next((m for m in messages if m.get("id") == rid), messages[0] if messages else None)
        if matched is None:
            raise SystemExit(f"no response for {method}")
        if "error" in matched:
            err = matched["error"]
            raise SystemExit(f"MCP error {err.get('code')}: {err.get('message')}")
        return matched.get("result")

    def initialize(self) -> dict[str, Any]:
        return self.send(
            "initialize",
            {
                "protocolVersion": "2024-11-05",
                "capabilities": {},
                "clientInfo": {"name": "demo-mcp-call.py", "version": "1.0.0"},
            },
        )

    def list_tools(self) -> list[dict[str, Any]]:
        result = self.send("tools/list")
        return result.get("tools", []) if isinstance(result, dict) else []

    def call_tool(self, name: str, args: dict[str, Any]) -> Any:
        return self.send("tools/call", {"name": name, "arguments": args})


def main() -> int:
    ap = argparse.ArgumentParser(description="Demo MCP request via McpToken bearer.")
    ap.add_argument("--url", default=os.environ.get("MCPGW_URL", "http://localhost:3200"),
                    help="Base URL of mcplocal (default: $MCPGW_URL or http://localhost:3200)")
    ap.add_argument("--project", default="sre",
                    help="Project name (default: sre). Must match the token's bound project.")
    ap.add_argument("--token", default=os.environ.get("MCPCTL_TOKEN"),
                    help="Raw mcpctl_pat_* bearer (default: $MCPCTL_TOKEN)")
    ap.add_argument("--tool", help="Optionally call a tool after tools/list")
    ap.add_argument("--args", default="{}", help="JSON-encoded arguments for --tool")
    ap.add_argument("--timeout", type=float, default=30.0)
    opts = ap.parse_args()

    if not opts.token:
        ap.error("--token or $MCPCTL_TOKEN required")

    endpoint = f"{opts.url.rstrip('/')}/projects/{opts.project}/mcp"
    print(f"→ POST {endpoint}")
    print(f"  Bearer: {opts.token[:16]}…")
    print()

    sess = McpSession(endpoint, bearer=opts.token, timeout=opts.timeout)

    info = sess.initialize()
    server_info = info.get("serverInfo", {}) if isinstance(info, dict) else {}
    print(f"initialize:  protocol={info.get('protocolVersion') if isinstance(info, dict) else '?'} "
          f"server={server_info.get('name', '?')}/{server_info.get('version', '?')} "
          f"sessionId={sess.session_id}")

    tools = sess.list_tools()
    print(f"tools/list:  {len(tools)} tool(s)")
    for t in tools:
        desc = (t.get("description") or "").splitlines()[0][:80]
        print(f"  - {t['name']}  {desc}")

    if opts.tool:
        try:
            args = json.loads(opts.args)
        except json.JSONDecodeError as e:
            raise SystemExit(f"--args must be valid JSON: {e}")
        print(f"\ntools/call: {opts.tool} {args}")
        result = sess.call_tool(opts.tool, args)
        print(json.dumps(result, indent=2)[:2000])

    return 0


if __name__ == "__main__":
    sys.exit(main())