#!/usr/bin/env python3
"""云巅实战营打卡 MCP 本地代理（只用 Python 标准库，Python 3.8+）。

给不能在后台完成 OAuth 浏览器登录的客户端（如 Kiro Crew）使用：
  1) 在终端登录一次（打开浏览器，用自己的邮箱 + 统一密码）：
       python3 ~/.daka/daka_mcp_proxy.py login
     只保存 refresh token 到 ~/.daka/token.json（权限 600），不保存密码。有效期 30 天。
  2) 在 MCP 配置里以 stdio 方式启动：
       python3 ~/.daka/daka_mcp_proxy.py serve
     代理自动刷新 access token，并把 MCP 请求转发到培训打卡 Gateway。

其他命令：status（查看登录状态）、logout（删除本地 token）。
"""
from __future__ import annotations

import base64
import hashlib
import json
import os
import secrets
import sys
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
import webbrowser
from http.server import BaseHTTPRequestHandler, HTTPServer

GATEWAY = "https://daka-kiro-cwi5rpnnoa.gateway.bedrock-agentcore.us-west-2.amazonaws.com/mcp"
DOMAIN = "https://daka-training-200226111370.auth.us-west-2.amazoncognito.com"
CLIENT_ID = "642n7p328iomd1t3paprf669qo"
# 依次尝试这些本机端口（均已在 Cognito 登记），避开被占用的端口
CALLBACK_PORTS = (53682, 53683, 53684, 8080)
TOKEN_FILE = os.path.join(os.path.expanduser("~"), ".daka", "token.json")
LOGIN_HINT = "请先在终端运行：python3 ~/.daka/daka_mcp_proxy.py login"


def log(msg: str) -> None:
    print(msg, file=sys.stderr, flush=True)


# ------------------------------------------------------------------ token store
def load_token() -> dict | None:
    try:
        with open(TOKEN_FILE, encoding="utf-8") as f:
            return json.load(f)
    except (OSError, ValueError):
        return None


def save_token(tok: dict) -> None:
    os.makedirs(os.path.dirname(TOKEN_FILE), mode=0o700, exist_ok=True)
    tmp = TOKEN_FILE + ".tmp"
    fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    with os.fdopen(fd, "w", encoding="utf-8") as f:
        json.dump(tok, f)
    os.replace(tmp, TOKEN_FILE)


def token_request(form: dict) -> dict:
    req = urllib.request.Request(
        f"{DOMAIN}/oauth2/token",
        data=urllib.parse.urlencode(form).encode(),
        headers={"content-type": "application/x-www-form-urlencoded"},
    )
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.loads(r.read())


def store_from_response(res: dict, old: dict | None = None) -> dict:
    tok = {
        "access_token": res["access_token"],
        "refresh_token": res.get("refresh_token") or (old or {}).get("refresh_token"),
        "expires_at": int(time.time()) + int(res.get("expires_in", 3600)),
        "login_at": (old or {}).get("login_at") or int(time.time()),
    }
    save_token(tok)
    return tok


_lock = threading.Lock()


def access_token(force_refresh: bool = False) -> str:
    with _lock:
        tok = load_token()
        if not tok or not tok.get("refresh_token"):
            raise RuntimeError("未登录。" + LOGIN_HINT)
        if force_refresh or time.time() > tok.get("expires_at", 0) - 120:
            try:
                res = token_request({"grant_type": "refresh_token", "client_id": CLIENT_ID, "refresh_token": tok["refresh_token"]})
            except urllib.error.HTTPError as e:
                raise RuntimeError(f"登录已过期（HTTP {e.code}）。" + LOGIN_HINT) from None
            tok = store_from_response(res, tok)
        return tok["access_token"]


# ------------------------------------------------------------------ login
class _CallbackHandler(BaseHTTPRequestHandler):
    state = ""
    result: dict = {}

    def do_GET(self):  # noqa: N802
        u = urllib.parse.urlparse(self.path)
        q = urllib.parse.parse_qs(u.query)
        ok = u.path == "/oauth/callback" and q.get("state", [""])[0] == self.state and "code" in q
        if ok:
            self.result["code"] = q["code"][0]
        body = ("登录成功，可以关闭此页面，回到终端。" if ok else "登录失败，请回到终端重试。").encode("utf-8")
        self.send_response(200 if ok else 400)
        self.send_header("content-type", "text/plain; charset=utf-8")
        self.send_header("content-length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, *a):
        pass


def login() -> None:
    srv = None
    for port in CALLBACK_PORTS:
        try:
            srv = HTTPServer(("127.0.0.1", port), _CallbackHandler)  # 只监听本机，登录完成即关闭
            break
        except OSError:
            continue
    if srv is None:
        sys.exit(f"本机端口 {', '.join(map(str, CALLBACK_PORTS))} 都被占用，请关闭占用的程序后重试。")
    redirect = f"http://127.0.0.1:{srv.server_address[1]}/oauth/callback"
    verifier = secrets.token_urlsafe(48)
    challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(b"=").decode()
    state = secrets.token_urlsafe(16)
    url = f"{DOMAIN}/oauth2/authorize?" + urllib.parse.urlencode({
        "response_type": "code", "client_id": CLIENT_ID, "redirect_uri": redirect,
        "scope": "openid email profile", "state": state,
        "code_challenge_method": "S256", "code_challenge": challenge,
    })
    _CallbackHandler.state = state
    _CallbackHandler.result = result = {}
    srv.timeout = 300
    print("正在打开浏览器登录页。如果没有自动打开，请手动访问：\n" + url, flush=True)
    webbrowser.open(url)
    deadline = time.time() + 300
    while "code" not in result and time.time() < deadline:
        srv.handle_request()
    srv.server_close()
    if "code" not in result:
        sys.exit("登录超时或失败，请重试。")
    res = token_request({"grant_type": "authorization_code", "client_id": CLIENT_ID, "code": result["code"], "redirect_uri": redirect, "code_verifier": verifier})
    store_from_response(res)
    print(f"登录成功，登录信息已保存到 {TOKEN_FILE}（30 天内有效）。")


def status() -> None:
    tok = load_token()
    if not tok:
        print("未登录。" + LOGIN_HINT)
        return
    try:
        access_token()
        days = 30 - (time.time() - tok.get("login_at", time.time())) / 86400
        print(f"已登录，约 {max(days, 0):.0f} 天后需要重新登录。")
    except RuntimeError as e:
        print(str(e))


# ------------------------------------------------------------------ stdio proxy
def post(body: bytes, session: str | None) -> tuple[int, dict, str]:
    for attempt in (0, 1):
        headers = {"content-type": "application/json", "accept": "application/json, text/event-stream",
                   "authorization": f"Bearer {access_token(force_refresh=attempt == 1)}"}
        if session:
            headers["mcp-session-id"] = session
        req = urllib.request.Request(GATEWAY, data=body, headers=headers, method="POST")
        try:
            with urllib.request.urlopen(req, timeout=120) as r:
                return r.status, {k.lower(): v for k, v in r.headers.items()}, r.read().decode("utf-8")
        except urllib.error.HTTPError as e:
            if e.code == 401 and attempt == 0:
                continue
            return e.code, {k.lower(): v for k, v in e.headers.items()}, e.read().decode("utf-8", "replace")
    raise AssertionError("unreachable")


def messages_from(text: str) -> list:
    text = text.strip()
    if not text:
        return []
    if text.startswith(("{", "[")):
        data = json.loads(text)
        return data if isinstance(data, list) else [data]
    out = []
    for line in text.splitlines():
        if line.startswith("data:") and line[5:].strip():
            out.append(json.loads(line[5:].strip()))
    return out


def write(msg: dict) -> None:
    sys.stdout.write(json.dumps(msg, ensure_ascii=False) + "\n")
    sys.stdout.flush()


def serve() -> None:
    session: str | None = None
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            msg = json.loads(line)
        except ValueError:
            continue
        rid = msg.get("id") if isinstance(msg, dict) else None
        try:
            status_code, headers, text = post(line.encode("utf-8"), session)
            session = headers.get("mcp-session-id", session)
            if status_code >= 400:
                raise RuntimeError(f"打卡服务返回 HTTP {status_code}")
            for m in messages_from(text):
                write(m)
        except Exception as e:  # noqa: BLE001
            log(f"daka proxy: {e}")
            if rid is not None:
                write({"jsonrpc": "2.0", "id": rid, "error": {"code": -32603, "message": str(e)}})


def main() -> None:
    cmd = sys.argv[1] if len(sys.argv) > 1 else "serve"
    if cmd == "login":
        login()
    elif cmd == "status":
        status()
    elif cmd == "logout":
        try:
            os.remove(TOKEN_FILE)
        except FileNotFoundError:
            pass
        print("已删除本地登录信息。")
    elif cmd == "serve":
        serve()
    else:
        sys.exit(__doc__)


if __name__ == "__main__":
    main()