#!/usr/bin/env python3
"""mekong-proxy.py — buyer-side local proxy for the Mekong x402 fleet.

Run this on YOUR OWN machine. Point any agent/framework at
http://127.0.0.1:8787/<slug> (e.g. /tokenprice, /brief, /websearch) and the
proxy forwards the call to https://api.mekongtradeintel.com, handling the
x402 payment flow for you:

  1. The call is forwarded as-is. If the endpoint serves it (free trial or
     already paid), the response passes straight back — no payment.
  2. If the endpoint answers HTTP 402, the proxy parses the
     PAYMENT-REQUIRED header, signs an EIP-712 TransferWithAuthorization
     with YOUR wallet, and retries once with the Payment-Signature header.
     This is the same payment core the fleet's code generator emits.

Your key is read from the X402_PAYER_KEY environment variable ONLY. It is
never stored, never logged, and never sent anywhere except as the signature
on your own payment. Without a key the proxy still forwards trial calls
and returns 402 responses untouched.

Spend control: a local daily cap (X402_DAILY_CAP_USD, default 1.00) is
enforced from a local ledger file (./mekong-proxy-ledger.json) BEFORE any
payment is signed — the same conservative design as the fleet spend guard:
every paid attempt counts at list price against today's (UTC) cap.

Usage:
  curl -O https://api.mekongtradeintel.com/tools/mekong-proxy.py
  pip install requests eth-account     # only for paid calls
  export X402_PAYER_KEY=0x...          # your Base wallet key (USDC on Base)
  python3 mekong-proxy.py              # 127.0.0.1:8787
  curl -X POST http://127.0.0.1:8787/tokenprice \
       -H 'Content-Type: application/json' -d '{"ticker":"ETH"}'

Framework use: register http://127.0.0.1:8787/<slug> as the tool/API URL in
LangChain, CrewAI, OpenAI SDK, or anything that can call an HTTP endpoint —
no x402 code in your framework at all. See
https://api.mekongtradeintel.com/tools/local-proxy.html

Env:
  X402_PAYER_KEY     payer wallet private key (optional; trial works without)
  MEKONG_PROXY_PORT  listen port (default 8787)
  MEKONG_PROXY_HOST  listen host (default 127.0.0.1 — keep it local!)
  X402_DAILY_CAP_USD daily spend cap in USD (default 1.00)
  MEKONG_API_BASE    fleet base URL (default https://api.mekongtradeintel.com)
  MEKONG_LEDGER      ledger path (default ./mekong-proxy-ledger.json)
"""
import base64
import json
import os
import sys
import time
import urllib.request
import urllib.error
from datetime import datetime, timezone
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

API_BASE = os.environ.get("MEKONG_API_BASE", "https://api.mekongtradeintel.com").rstrip("/")
HOST = os.environ.get("MEKONG_PROXY_HOST", "127.0.0.1")
PORT = int(os.environ.get("MEKONG_PROXY_PORT", "8787"))
CAP_USD = float(os.environ.get("X402_DAILY_CAP_USD", "1.00"))
LEDGER_PATH = os.environ.get("MEKONG_LEDGER", "./mekong-proxy-ledger.json")
PAYER_KEY = os.environ.get("X402_PAYER_KEY", "").strip()

_account = None
if PAYER_KEY:
    try:
        from eth_account import Account
        _account = Account.from_key(PAYER_KEY)
    except ImportError:
        print("mekong-proxy: X402_PAYER_KEY is set but eth-account is not installed.\n"
              "  pip install eth-account   (trial calls still work; paid calls will return 402 untouched)",
              file=sys.stderr)
    except Exception as e:
        print(f"mekong-proxy: could not load X402_PAYER_KEY ({e}); running trial-only.", file=sys.stderr)

_catalog = {"at": 0.0, "endpoints": {}}


def load_catalog():
    """Fleet catalogue from the clinic codegen (live prices + URLs), 5-min cache."""
    now = time.time()
    if _catalog["endpoints"] and now - _catalog["at"] < 300:
        return _catalog["endpoints"]
    try:
        req = urllib.request.Request(API_BASE + "/clinic/codegen",
                                     headers={"User-Agent": "mekong-proxy/1.0"})
        with urllib.request.urlopen(req, timeout=15) as r:
            data = json.loads(r.read().decode("utf-8"))
        eps = {}
        for e in data.get("endpoints", []):
            eps[e["slug"]] = {"method": e.get("method", "POST"),
                              "url": e["url"], "price_usd": float(e.get("price_usd", 0))}
        if eps:
            _catalog["endpoints"] = eps
            _catalog["at"] = now
    except Exception as e:
        print(f"mekong-proxy: catalogue fetch failed ({e}); keeping last cache.", file=sys.stderr)
    return _catalog["endpoints"]


def _ledger():
    try:
        with open(LEDGER_PATH) as f:
            return json.load(f)
    except Exception:
        return {}


def _today():
    return datetime.now(timezone.utc).strftime("%Y-%m-%d")


def spent_today(slug):
    wallet = _account.address if _account else "nokey"
    return _ledger().get(f"{wallet}|{slug}|{_today()}", {}).get("usd", 0.0)


def guard_check(slug, price_usd):
    spent = spent_today(slug)
    if spent + price_usd > CAP_USD:
        raise RuntimeError(
            f"SPEND CAP BLOCKED: {slug} costs ${price_usd} but only "
            f"${max(0.0, CAP_USD - spent):.4f} of today's ${CAP_USD} cap remains "
            f"(already spent ${spent} today, UTC). Raise X402_DAILY_CAP_USD or wait until tomorrow.")


def guard_record(slug, price_usd):
    db = _ledger()
    wallet = _account.address if _account else "nokey"
    k = f"{wallet}|{slug}|{_today()}"
    row = db.get(k, {"usd": 0.0, "calls": 0})
    row["usd"] = round(row["usd"] + price_usd, 6)
    row["calls"] = row["calls"] + 1
    row["updated"] = datetime.now(timezone.utc).isoformat()
    db[k] = row
    with open(LEDGER_PATH, "w") as f:
        json.dump(db, f, indent=2)


def http_call(method, url, body_bytes, headers, timeout=60):
    """One HTTP call via urllib. Returns (status, headers_dict, body_bytes)."""
    req = urllib.request.Request(url, data=body_bytes if method != "GET" else None,
                                 method=method)
    for k, v in headers.items():
        req.add_header(k, v)
    req.add_header("User-Agent", "mekong-proxy/1.0")
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return r.status, dict(r.headers.items()), r.read()
    except urllib.error.HTTPError as e:
        return e.code, dict(e.headers.items() if e.headers else {}), e.read()
    except Exception as e:
        return None, {}, (json.dumps({"proxy_error": str(e)}).encode())


def sign_payment(req402):
    """Sign an x402 exact payment. Same core as the clinic codegen snippets."""
    import secrets
    from eth_account.messages import encode_typed_data
    now = int(time.time())
    auth = {"from": _account.address, "to": req402["payTo"], "value": str(req402["amount"]),
            "validAfter": str(now - 60),
            "validBefore": str(now + int(req402.get("maxTimeoutSeconds", 300))),
            "nonce": "0x" + secrets.token_hex(32)}
    typed = {"types": {"EIP712Domain": [{"name": "name", "type": "string"}, {"name": "version", "type": "string"}, {"name": "chainId", "type": "uint256"}, {"name": "verifyingContract", "type": "address"}],
                       "TransferWithAuthorization": [{"name": "from", "type": "address"}, {"name": "to", "type": "address"}, {"name": "value", "type": "uint256"}, {"name": "validAfter", "type": "uint256"}, {"name": "validBefore", "type": "uint256"}, {"name": "nonce", "type": "bytes32"}]},
             "primaryType": "TransferWithAuthorization",
             "domain": {"name": "USD Coin", "version": "2", "chainId": 8453, "verifyingContract": req402["asset"]},
             "message": {**auth, "value": int(auth["value"]), "validAfter": int(auth["validAfter"]), "validBefore": int(auth["validBefore"])}}
    sig = _account.sign_message(encode_typed_data(full_message=typed)).signature.hex()
    if not sig.startswith("0x"):
        sig = "0x" + sig
    payload = {"x402Version": 2, "scheme": "exact", "network": req402["network"],
               "payload": {"signature": sig, "authorization": auth}, "accepted": req402}
    return base64.b64encode(json.dumps(payload).encode()).decode()


def forward(slug, method, body_bytes, content_type):
    """Forward one call to the fleet, paying on 402 when a key is configured."""
    eps = load_catalog()
    ep = eps.get(slug)
    if not ep:
        return 404, {"Content-Type": "application/json"}, json.dumps(
            {"error": f"unknown endpoint slug '{slug}'", "hint": "GET /catalogue lists the fleet slugs"}).encode()
    url = ep["url"]
    headers = {}
    if content_type:
        headers["Content-Type"] = content_type
    status, hdrs, body = http_call(method, url, body_bytes, headers)
    if status != 402:
        return status, hdrs, body
    # --- 402: payment required ---
    if _account is None:
        out = dict(hdrs)
        out["X-Mekong-Proxy"] = "no payer key configured — 402 returned untouched (set X402_PAYER_KEY to pay automatically)"
        return status, out, body
    pr = hdrs.get("PAYMENT-REQUIRED") or hdrs.get("Payment-Required") or hdrs.get("payment-required")
    if not pr:
        return status, hdrs, body
    try:
        req402 = json.loads(base64.b64decode(pr))["accepts"][0]
    except Exception:
        return status, hdrs, body
    price_usd = int(req402["amount"]) / 1e6
    try:
        guard_check(slug, price_usd)
    except RuntimeError as e:
        return 402, {"Content-Type": "application/json"}, json.dumps(
            {"error": str(e), "endpoint": slug, "price_usd": price_usd}).encode()
    header = sign_payment(req402)
    headers2 = dict(headers)
    headers2["Payment-Signature"] = header
    status2, hdrs2, body2 = http_call(method, url, body_bytes, headers2)
    if status2 != 402:
        guard_record(slug, price_usd)
    out = dict(hdrs2)
    out["X-Mekong-Proxy"] = f"paid ${price_usd} for {slug} via x402 (payer {_account.address})"
    return status2, out, body2


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log_message(self, fmt, *args):
        sys.stderr.write("mekong-proxy: %s\n" % (fmt % args))

    def _send(self, status, headers, body):
        if status is None:
            status = 502
        self.send_response(status)
        ctype = headers.get("Content-Type") or headers.get("content-type") or "application/json"
        self.send_header("Content-Type", ctype)
        for k in ("X-Trial-Remaining", "PAYMENT-RESPONSE", "X-Permalink", "X-Mekong-Proxy"):
            for hk, hv in headers.items():
                if hk.lower() == k.lower():
                    self.send_header(k, hv)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        if self.command != "HEAD":
            self.wfile.write(body)

    def _handle(self):
        path = self.path.split("?")[0].strip("/")
        if path in ("", "health"):
            wallet = _account.address if _account else None
            self._send(200, {"Content-Type": "application/json"}, json.dumps({
                "ok": True, "proxy": "mekong-proxy/1.0", "payer": wallet,
                "paid_calls_enabled": _account is not None,
                "daily_cap_usd": CAP_USD, "api_base": API_BASE,
                "catalogue_size": len(load_catalog()),
                "usage": "GET|POST /<slug> forwards to the fleet (e.g. /tokenprice, /brief); GET /catalogue lists slugs",
            }).encode())
            return
        if path == "catalogue":
            self._send(200, {"Content-Type": "application/json"},
                       json.dumps({"endpoints": load_catalog()}, indent=2).encode())
            return
        slug = path.split("/")[0].lower()
        n = int(self.headers.get("Content-Length", 0) or 0)
        body = self.rfile.read(min(n, 32 * 1024 * 1024)) if n > 0 else (b"" if self.command in ("POST", "PUT", "PATCH") else None)
        t0 = time.time()
        status, hdrs, resp = forward(slug, self.command, body if body is not None else b"",
                                     self.headers.get("Content-Type"))
        print(f"mekong-proxy: {self.command} /{slug} -> {status} in {int((time.time()-t0)*1000)}ms", flush=True)
        self._send(status, hdrs, resp)

    do_GET = _handle
    do_POST = _handle
    do_PUT = _handle


if __name__ == "__main__":
    eps = load_catalog()
    server = ThreadingHTTPServer((HOST, PORT), Handler)  # binds here; raises if the port is taken
    print(f"mekong-proxy listening on http://{HOST}:{PORT} — {len(eps)} fleet endpoints, "
          f"payer {_account.address if _account else '(none — trial-only mode)'}, daily cap ${CAP_USD}", flush=True)
    server.serve_forever()