#!/usr/bin/env python3
"""Local stdio bridge between the Claude desktop app and Amazon's MCP servers.

Neither of Amazon's MCP servers can be added to Claude as a remote connector:
the OAuth handshake asks for a scope Amazon's login server will not grant to a
generic client. This script sidesteps that by holding credentials for an
application Amazon has already approved, minting the access token itself, and
forwarding messages between Claude and Amazon.

    python3 mcp_proxy.py --server ads   --readonly --selftest
    python3 mcp_proxy.py --server spapi --selftest

Credentials are read from .env.ads / .env.spapi in the same folder as this file.
Nothing is uploaded anywhere; the token exchange goes straight to Amazon.
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
LWA_TOKEN_URL = "https://api.amazon.com/auth/o2/token"

SERVERS = {
    "ads": {
        "url": "https://advertising-ai.amazon.com/mcp",
        "protocol": "2025-06-18",
        "env_file": ".env.ads",
        "keys": ("ADS_CLIENT_ID", "ADS_CLIENT_SECRET", "ADS_REFRESH_TOKEN"),
        "client_id_header": True,
    },
    "spapi": {
        "url": "https://sellingpartner-ai.amazon.com/mcp",
        "protocol": "2025-03-26",
        "env_file": ".env.spapi",
        "keys": ("SPAPI_LWA_CLIENT_ID", "SPAPI_LWA_CLIENT_SECRET", "SPAPI_REFRESH_TOKEN"),
        "client_id_header": False,
    },
}


def log(msg):
    """stderr only. stdout is the MCP channel; anything else there breaks it."""
    print(f"[mcp-proxy] {msg}", file=sys.stderr, flush=True)


def read_env(filename):
    path = os.path.join(HERE, filename)
    if not os.path.exists(path):
        sys.exit(f"error: {path} not found. Create it with your Amazon credentials.")
    values = {}
    with open(path) as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith("#") and "=" in line:
                k, v = line.split("=", 1)
                values[k.strip()] = v.strip().strip("'\"")
    return values


class Upstream:
    def __init__(self, name, readonly=False):
        self.cfg = SERVERS[name]
        self.name = name
        self.readonly = readonly and name == "ads"
        env = read_env(self.cfg["env_file"])
        cid, secret, refresh = self.cfg["keys"]
        missing = [k for k in (cid, secret, refresh) if k not in env]
        if missing:
            sys.exit(f"error: {self.cfg['env_file']} is missing {', '.join(missing)}")
        self.client_id, self.secret, self.refresh = env[cid], env[secret], env[refresh]
        self.token = None
        self.expires = 0
        self.session = None
        self.readonly_names = None

    def access_token(self):
        if self.token and time.time() < self.expires - 120:
            return self.token
        body = urllib.parse.urlencode({
            "grant_type": "refresh_token", "refresh_token": self.refresh,
            "client_id": self.client_id, "client_secret": self.secret}).encode()
        req = urllib.request.Request(
            LWA_TOKEN_URL, data=body,
            headers={"Content-Type": "application/x-www-form-urlencoded"})
        try:
            with urllib.request.urlopen(req, timeout=30) as r:
                payload = json.load(r)
        except urllib.error.HTTPError as e:
            detail = json.loads(e.read() or b"{}").get("error_description", f"HTTP {e.code}")
            sys.exit(f"error: Amazon rejected the credentials — {detail}")
        self.token = payload["access_token"]
        self.expires = time.time() + int(payload.get("expires_in", 3600))
        log(f"{self.name}: access token refreshed")
        return self.token

    def forward(self, message):
        headers = {"Authorization": f"Bearer {self.access_token()}",
                   "Content-Type": "application/json",
                   "Accept": "application/json, text/event-stream"}
        if self.cfg["client_id_header"]:
            headers["Amazon-Advertising-API-ClientId"] = self.client_id
        if self.session:
            headers["Mcp-Session-Id"] = self.session
        req = urllib.request.Request(self.cfg["url"], data=json.dumps(message).encode(),
                                     method="POST", headers=headers)
        try:
            with urllib.request.urlopen(req, timeout=120) as r:
                sid = r.headers.get("Mcp-Session-Id")
                if sid:
                    self.session = sid
                raw = r.read().decode()
                if not raw:
                    return None
                # Streamable HTTP may answer as SSE; take the last data: line.
                if raw.lstrip().startswith(("event:", "data:")):
                    data = [l[5:].strip() for l in raw.splitlines() if l.startswith("data:")]
                    raw = data[-1] if data else "{}"
                return json.loads(raw)
        except urllib.error.HTTPError as e:
            detail = e.read().decode()[:300]
            log(f"{self.name}: upstream HTTP {e.code}: {detail}")
            if message.get("id") is None:
                return None
            return {"jsonrpc": "2.0", "id": message["id"],
                    "error": {"code": -32000, "message": f"upstream HTTP {e.code}",
                              "data": detail}}

    def filter_list(self, message, response):
        """Sort the catalogue, and in readonly mode drop everything that writes.

        Sorting is not cosmetic. Tool definitions sit at the front of the prompt
        and prompt caching is a byte-exact prefix match, but Amazon returns the
        same tools in a different order on each connection. Unsorted, every
        reconnect writes a fresh cache entry instead of reading the old one.
        """
        if not response or "result" not in response:
            return response
        if message.get("method") != "tools/list":
            return response
        tools = sorted(response["result"].get("tools", []), key=lambda t: t["name"])
        if self.readonly:
            kept = [t for t in tools if (t.get("annotations") or {}).get("readOnlyHint") is True]
            self.readonly_names = {t["name"] for t in kept}
            log(f"readonly: exposing {len(kept)} of {len(tools)} tools")
            tools = kept
        response["result"]["tools"] = tools
        return response

    def guard_call(self, message):
        """Refuse tools/call for anything outside the read-only set."""
        if not self.readonly or message.get("method") != "tools/call":
            return None
        name = (message.get("params") or {}).get("name")
        if self.readonly_names is None:
            return {"jsonrpc": "2.0", "id": message.get("id"),
                    "error": {"code": -32001,
                              "message": "call tools/list first"}}
        if name not in self.readonly_names:
            log(f"BLOCKED {name}")
            return {"jsonrpc": "2.0", "id": message.get("id"),
                    "error": {"code": -32001,
                              "message": f"'{name}' is not read-only. Blocked by --readonly."}}
        return None


def selftest(up):
    init = up.forward({"jsonrpc": "2.0", "id": 1, "method": "initialize",
                       "params": {"protocolVersion": up.cfg["protocol"], "capabilities": {},
                                  "clientInfo": {"name": "selftest", "version": "1.0"}}})
    info = (init or {}).get("result", {}).get("serverInfo", {})
    print(f"upstream : {info.get('name')} v{info.get('version')}")
    up.forward({"jsonrpc": "2.0", "method": "notifications/initialized"})
    tl = up.forward({"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}})
    tools = (tl or {}).get("result", {}).get("tools", [])
    ro = [t for t in tools if (t.get("annotations") or {}).get("readOnlyHint") is True]
    print(f"tools    : {len(tools)} total" + (f", {len(ro)} read-only" if ro else ""))
    print(f"mode     : {'READ-ONLY' if up.readonly else 'FULL'}")
    print("selftest ok")


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--server", choices=sorted(SERVERS), default="ads")
    ap.add_argument("--readonly", action="store_true",
                    help="Ads only: expose just the tools Amazon marks read-only")
    ap.add_argument("--selftest", action="store_true",
                    help="check the connection and exit")
    args = ap.parse_args()

    up = Upstream(args.server, args.readonly)
    if args.selftest:
        return selftest(up)

    log(f"started server={args.server} mode={'readonly' if up.readonly else 'full'}")
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            message = json.loads(line)
        except json.JSONDecodeError:
            continue
        blocked = up.guard_call(message)
        if blocked:
            sys.stdout.write(json.dumps(blocked) + "\n")
            sys.stdout.flush()
            continue
        response = up.forward(message)
        if response is None:
            continue
        response = up.filter_list(message, response)
        sys.stdout.write(json.dumps(response) + "\n")
        sys.stdout.flush()


if __name__ == "__main__":
    main()
