"""
Put Squidbrake in front of ANY app's MCP connector: Stripe, GitHub, Slack, Gmail, Linear, Notion,
a database, an internal tool... The agent sees the app's normal tools; every call is checked against
the rules, recorded, and held for a human when a rule says so, before it reaches the app.

    python gateway_proxy.py --app stripe -- npx -y @stripe/mcp --tools=all
    python gateway_proxy.py --app github -- docker run -i --rm -e GITHUB_PERSONAL_ACCESS_TOKEN ghcr.io/github/github-mcp-server
    python gateway_proxy.py --app acme   -- python demo_apps_mcp.py        (the sandbox company)
    python gateway_proxy.py --app crm --url https://crm.example.com/mcp    (a remote MCP server)

Served over HTTP, for agents that connect to MCP by URL (ChatGPT and claude.ai connectors, Devin, n8n, cloud agents):

    python gateway_proxy.py --app github --serve 0.0.0.0:9000 --token SECRET -- npx -y @modelcontextprotocol/server-github
    -> the agent connects to http(s)://this-host:9000/mcp with "Authorization: Bearer SECRET" (or ?token=SECRET)

Calls show up in the dashboard as <app>.<tool>, e.g. stripe.create_refund, so rules can target them.
The app's own credentials (e.g. STRIPE_SECRET_KEY) go in the agent's MCP config env as usual; they are
passed to the app, never to the gateway. Gateway settings: see gw_async.py.
Easiest setup:  python connect.py wrap --app NAME --agent claude-code -- <the app's MCP command>
"""
from __future__ import annotations

import argparse
import asyncio
import json
import os
import re
import sys
import time
import uuid
from contextlib import asynccontextmanager

from mcp import Client, types
from mcp.client.stdio import StdioServerParameters
from mcp.server.lowlevel import Server
from mcp.server.stdio import stdio_server

import gw_async as gw
from gw_async import log

APP = "app"
_upstream: Client | None = None

def gateway_tools() -> list[types.Tool]:
    """The gateway's own tools, named after the app (acme_check_approval) so they never clash with another
    connector's in agents that keep one list of tools."""
    return [
        types.Tool(
            name=gw.CHECK_TOOL,
            description=f"Finish a {APP} call Squidbrake was holding for human approval (WAITING FOR HUMAN "
                        "APPROVAL): runs it if approved, reports if rejected, or keeps waiting a little longer.",
            input_schema={"type": "object", "properties": {"event_id": {"type": "string"}}, "required": ["event_id"]},
        ),
        types.Tool(name=gw.DECISIONS_TOOL, description=gw.DECISIONS_HELP,
                   input_schema={"type": "object", "properties": {}}),
        types.Tool(name=gw.ASK_TOOL, description=gw.ASK_HELP, input_schema=gw.ASK_SCHEMA),
    ]


_justified: set[str] = set()     # the app's tools given the extra argument (an app's own `justification` stays its)


def with_justification(tool: types.Tool) -> types.Tool:
    """The app's tool, plus an optional `justification` argument (rules.yaml `justification:`). It goes to Squidbrake
    only: call_tool takes it out before the app sees the call."""
    schema = dict(tool.input_schema or {})
    props = schema.get("properties")
    if schema.get("type", "object") != "object" or not isinstance(props, dict) or "justification" in props:
        return tool
    _justified.add(tool.name)
    return tool.model_copy(update={"input_schema": {**schema, "properties": {**props,
                                                                             "justification": gw.JUSTIFICATION_PARAM}}})


def replay(call: dict):
    async def run():
        return await _upstream.call_tool(call["tool"], call.get("args") or {})
    return run


def text_result(message: str, is_error: bool = False) -> types.CallToolResult:
    return types.CallToolResult(content=[types.TextContent(type="text", text=message)], is_error=is_error)


def summarize(result: types.CallToolResult) -> dict:
    """What the gateway records about the app's answer (text is kept; images/files are noted, not stored)."""
    parts = []
    for c in result.content or []:
        parts.append(c.text if getattr(c, "type", "") == "text" else f"[{getattr(c, 'type', 'content')}]")
    out: dict = {"content": parts[0] if len(parts) == 1 else parts}
    if result.structured_content is not None:
        out["structured"] = result.structured_content
    return out


async def report(event_id: str, outcome) -> types.CallToolResult:
    status, value, ms = outcome
    if status == "error":  # couldn't reach the app, or it crashed
        await gw.gw_result(event_id, error=value, duration_ms=ms)
        return text_result(f"ERROR from {APP}: {value}", is_error=True)
    summary = summarize(value)
    err = None
    if value.is_error:
        err = summary["content"] if isinstance(summary["content"], str) else json.dumps(summary["content"])[:2000]
    await gw.gw_result(event_id, output=summary, error=err, duration_ms=ms)
    return value


async def list_tools(ctx, params) -> types.ListToolsResult:
    result = await _upstream.list_tools(cursor=params.cursor if params else None)
    tools = [with_justification(t) for t in result.tools]
    if not (params and params.cursor):
        tools += gateway_tools()
    return types.ListToolsResult(tools=tools, next_cursor=result.next_cursor)


IDLE_NEW_CONVERSATION = 1800       # seconds without a call before the same client starts a new conversation
_recent: dict[tuple, tuple[str, float]] = {}


def _conversation(ctx) -> str | None:
    """Served over HTTP, which conversation a call belongs to, so each one's chain is checked on its own:
    the X-Squidbrake-Session header if the caller sets one (n8n, scripts), else the MCP session of the connection,
    else (the 2026 protocol has no sessions) the same client's calls until it goes quiet for 30 minutes.
    None over stdio: one process is one conversation."""
    request = getattr(ctx, "request", None)
    if request is None or not hasattr(request, "headers"):
        return None
    headers = request.headers
    if given := headers.get("x-squidbrake-session"):
        return f"{gw.source()}-{re.sub(r'[^A-Za-z0-9_.:-]', '', given)[:60]}"
    conn = getattr(getattr(ctx, "session", None), "_connection", None)
    if sid := getattr(conn, "session_id", None) or headers.get("mcp-session-id"):
        return f"{gw.source()}-{sid[:24]}"
    params = getattr(conn, "client_params", None)
    info = getattr(params, "client_info", None) or getattr(params, "clientInfo", None)
    who = (gw.source(), getattr(info, "name", None) or "client", getattr(getattr(request, "client", None), "host", None))
    now = time.monotonic()
    sid, last = _recent.get(who, (None, 0.0))
    if not sid or now - last > IDLE_NEW_CONVERSATION:
        sid = f"{gw.source()}-{uuid.uuid4().hex[:12]}"
    _recent[who] = (sid, now)
    return sid


def _caller(ctx) -> None:
    """Served over HTTP: a caller that sends its own agent key (X-Gateway-Key) is checked and recorded as itself, and
    may say what to call it (X-Squidbrake-Source). The gateway's /mcp/<name> endpoint sends both."""
    request = getattr(ctx, "request", None)
    headers = getattr(request, "headers", None)
    if headers is None:
        return
    if key := headers.get("x-gateway-key"):
        gw.CURRENT_KEY.set(key)
    if name := re.sub(r"[^A-Za-z0-9_.:@-]", "", headers.get("x-squidbrake-source") or "")[:60]:
        gw.CURRENT_SOURCE.set(name)


async def call_tool(ctx, params: types.CallToolRequestParams) -> types.CallToolResult:
    _caller(ctx)
    if conv := _conversation(ctx):
        gw.CURRENT_SESSION.set(conv)
    args = params.arguments or {}
    on_text = lambda t: text_result(t, is_error=t.startswith("NOT RUN") or t.startswith("Unknown"))
    if params.name == gw.DECISIONS_TOOL:
        return text_result(await gw.recent_decisions())
    if params.name == gw.CHECK_TOOL:
        return await gw.check_approval(str(args.get("event_id", "")), on_done=report, on_text=on_text)
    if params.name == gw.ASK_TOOL:
        return text_result(await gw.ask_user(args.get("question"), args.get("event_id")))

    args = dict(args)
    why = args.pop("justification", None) if params.name in _justified else None   # for Squidbrake and the approver; the app never sees it
    call = {"tool": params.name, "args": args}
    return await gw.guard(f"{APP}.{params.name}", args, replay(call), on_done=report, on_text=on_text, kind="mcp",
                          call=call, metadata={"justification": str(why)[:2000]} if why else None)


def upstream_target(args):
    if args.url:
        if not args.header and not args.oauth_store:
            return args.url
        # a remote server that needs its own auth (Authorization: Bearer ...): ${VARS} come from the environment,
        # so the secret can stay in the MCP config's env instead of its arguments
        from mcp.client.streamable_http import create_mcp_http_client, streamable_http_client
        headers = {}
        for h in args.header:
            k, _, v = h.partition(":")
            headers[k.strip()] = os.path.expandvars(v.strip())
        auth = None
        if args.oauth_store:     # signed in through the dashboard (mcp_oauth.py): use the tokens, refresh them in time
            import mcp_oauth
            auth = mcp_oauth.provider(args.url, mcp_oauth.FileTokenStorage(args.oauth_store),
                                      redirect_uri=os.getenv("SQUIDBRAKE_OAUTH_REDIRECT", "http://localhost/unused"))
        return streamable_http_client(args.url, http_client=create_mcp_http_client(headers=headers or None, auth=auth))
    if not args.command:
        sys.exit("give the app's MCP command after --, or --url for a remote server")
    return StdioServerParameters(command=args.command[0], args=args.command[1:], env=dict(os.environ))


@asynccontextmanager
async def lifespan(server, target):
    global _upstream
    async with Client(target) as upstream:
        _upstream = upstream
        tools = (await upstream.list_tools()).tools
        for t in tools:              # known before the first call too (an agent may reuse a list it got earlier)
            with_justification(t)
        log(f"guarding '{APP}' ({len(tools)} tools) -> gateway {gw.GATEWAY_URL} as '{gw.SOURCE}' (session {gw.SESSION})")
        yield {}


def make_server(target) -> Server:
    return Server(
        f"gateway-{APP}",
        instructions=f"{APP} tools. " + gw.agent_rules(),
        on_list_tools=list_tools, on_call_tool=call_tool,
        lifespan=lambda s: lifespan(s, target),
    )


async def amain(args) -> None:
    server = make_server(upstream_target(args))
    async with stdio_server() as (read, write):
        await server.run(read, write, server.create_initialization_options())


# --------------------------------------------------------------------------- served over HTTP (--serve)
# For agents that connect to MCP by URL rather than by starting a command: ChatGPT and claude.ai connectors,
# Devin, n8n, cloud agents. One process guards one app; every connected conversation is checked on its own.

def require_token(app, token: str):
    """ASGI wrapper: only callers with the token reach the MCP endpoint. It can come as `Authorization: Bearer`,
    `X-Squidbrake-Token`, or `?token=` for clients that take only a URL."""
    import hmac
    from urllib.parse import parse_qs

    async def guarded(scope, receive, send):
        if scope["type"] != "http":
            return await app(scope, receive, send)
        headers = {k.decode("latin-1").lower(): v.decode("latin-1") for k, v in scope.get("headers") or []}
        given = headers.get("x-squidbrake-token") or ""
        if headers.get("authorization", "").lower().startswith("bearer "):
            given = headers["authorization"][7:].strip()
        given = given or (parse_qs(scope.get("query_string", b"").decode("latin-1")).get("token") or [""])[0]
        if not hmac.compare_digest(given.encode(), token.encode()):
            body = b'{"error": "missing or wrong Squidbrake token"}'
            await send({"type": "http.response.start", "status": 401,
                        "headers": [(b"content-type", b"application/json"), (b"www-authenticate", b"Bearer")]})
            return await send({"type": "http.response.body", "body": body})
        return await app(scope, receive, send)
    return guarded


def _alive(pid: int) -> bool:
    if sys.platform == "win32":     # never os.kill here: on Windows it terminates the process
        import ctypes
        k32 = ctypes.windll.kernel32
        h = k32.OpenProcess(0x1000, False, pid)          # PROCESS_QUERY_LIMITED_INFORMATION
        if not h:
            return False
        code = ctypes.c_ulong()
        k32.GetExitCodeProcess(h, ctypes.byref(code))
        k32.CloseHandle(h)
        return code.value == 259                         # STILL_ACTIVE
    try:
        os.kill(pid, 0)
        return True
    except ProcessLookupError:
        return False
    except PermissionError:
        return True


def _exit_with_parent() -> None:
    """Started by the gateway (mcp_hub.py): go when it goes, even if it crashed and couldn't stop us."""
    pid = int(os.getenv("SQUIDBRAKE_PARENT_PID") or 0)
    if not pid:
        return

    def watch():
        while _alive(pid):
            time.sleep(2)
        log("the gateway that started this proxy is gone; stopping")
        os._exit(0)
    import threading
    threading.Thread(target=watch, daemon=True).start()


def serve(args) -> None:
    import uvicorn
    _exit_with_parent()
    host, _, port = args.serve.rpartition(":")
    host = host or "127.0.0.1"
    token = args.token or os.getenv("SQUIDBRAKE_PROXY_TOKEN", "")
    local = host in ("127.0.0.1", "localhost", "::1")
    if not token and not local:
        sys.exit("--serve on a public address needs a token: --token SECRET or SQUIDBRAKE_PROXY_TOKEN, so only your "
                 "agents can call it")
    server = make_server(upstream_target(args))
    app = server.streamable_http_app(streamable_http_path=args.path, host=host)
    if token:
        app = require_token(app, token)
    log(f"serving '{APP}' over MCP at http://{host}:{port}{args.path}" + (" (token required)" if token else ""))
    uvicorn.run(app, host=host, port=int(port), log_level="warning")


def main() -> None:
    global APP
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--app", required=True, help="short name used in the dashboard and rules, e.g. stripe")
    p.add_argument("--url", help="a remote (streamable HTTP) MCP server instead of a command")
    p.add_argument("--header", action="append", default=[], metavar="'NAME: VALUE'",
                   help="with --url: a header the remote server needs, e.g. 'Authorization: Bearer ${GITHUB_TOKEN}' "
                        "(${VARS} are read from the environment); repeatable")
    p.add_argument("--oauth-store", metavar="FILE",
                   help="with --url: sign in with the OAuth tokens in FILE (written by the dashboard's Connect, "
                        "mcp_oauth.py) and refresh them as they run out")
    p.add_argument("--serve", metavar="HOST:PORT",
                   help="serve over HTTP instead of stdio, for agents that connect by URL (ChatGPT, claude.ai, Devin, "
                        "n8n...), e.g. --serve 0.0.0.0:9000")
    p.add_argument("--path", default="/mcp", help="with --serve: the endpoint's path (default /mcp)")
    p.add_argument("--token", help="with --serve: callers must send this (Bearer, X-Squidbrake-Token or ?token=); "
                                   "or set SQUIDBRAKE_PROXY_TOKEN. Required on a public address")
    p.add_argument("command", nargs=argparse.REMAINDER, help="-- then the app's MCP server command")
    args = p.parse_args()
    if args.command and args.command[0] == "--":
        args.command = args.command[1:]
    APP = args.app
    gw.set_tool_prefix(APP)
    gw.set_replay(replay)
    gw.cleanup_pending()
    import logging
    logging.getLogger("httpx").setLevel(logging.WARNING)
    if not gw.GATEWAY_API_KEY:
        log("warning: GATEWAY_API_KEY is not set; the gateway will reject calls unless auth is off")
    if args.serve:
        serve(args)
    else:
        asyncio.run(amain(args))


if __name__ == "__main__":
    main()