#!/usr/bin/env python3
"""Bearer token for the Superset MCP server (https://superset.swiftdrive-tech.ru/mcp).

Keycloak realm `infra`, public client `superset-mcp`, Device Authorization Grant.
First run opens a browser login (realm 2FA applies); after that the cached
refresh token renews silently until the Keycloak session ends.

  mcp-token.py            print the access token
  mcp-token.py --header   print {"Authorization": "Bearer ..."} (Claude Code headersHelper)
  mcp-token.py --logout   drop the cached token
"""
import json
import os
import subprocess
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

ISSUER = os.environ.get("SUPERSET_MCP_ISSUER", "https://auth.swiftdrive-tech.ru/realms/infra")
CLIENT_ID = os.environ.get("SUPERSET_MCP_CLIENT_ID", "superset-mcp")
CACHE = os.path.expanduser("~/.cache/superset-mcp/token.json")
MIN_TTL = 300  # renew when fewer seconds than this remain


def post(path, **form):
    body = urllib.parse.urlencode({"client_id": CLIENT_ID, **form}).encode()
    req = urllib.request.Request(f"{ISSUER}/protocol/openid-connect/{path}", body)
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            return json.load(resp)
    except urllib.error.HTTPError as e:
        try:
            return json.load(e)
        except ValueError:
            return {"error": f"HTTP {e.code}"}
    except urllib.error.URLError as e:
        sys.exit(f"superset-mcp: {ISSUER} unreachable: {e.reason}")


def save(tok):
    now = time.time()
    tok["expires_at"] = now + tok["expires_in"]
    tok["refresh_expires_at"] = now + tok.get("refresh_expires_in", 0)
    cache_dir = os.path.dirname(CACHE)
    os.makedirs(cache_dir, mode=0o700, exist_ok=True)
    os.chmod(cache_dir, 0o700)
    tmp = f"{CACHE}.{os.getpid()}"
    fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    with os.fdopen(fd, "w") as f:
        json.dump(tok, f)
    os.replace(tmp, CACHE)  # concurrent headersHelper runs never see a torn file
    return tok


def load():
    try:
        with open(CACHE) as f:
            return json.load(f)
    except (OSError, ValueError):
        return None


def device_login():
    dev = post("auth/device")
    if "device_code" not in dev:
        sys.exit(f"device authorization failed: {dev}")
    url = dev.get("verification_uri_complete") or dev["verification_uri"]
    print(f"Superset MCP login: open {url} (code {dev['user_code']})", file=sys.stderr)
    if sys.platform == "darwin":
        subprocess.run(["open", url], check=False)
    interval = dev.get("interval", 5)
    deadline = time.time() + dev.get("expires_in", 600)
    while time.time() < deadline:
        time.sleep(interval)
        tok = post("token", grant_type="urn:ietf:params:oauth:grant-type:device_code",
                   device_code=dev["device_code"])
        if "access_token" in tok:
            return save(tok)
        err = tok.get("error")
        if err == "slow_down":
            interval += 5
        elif err != "authorization_pending":
            sys.exit(f"login failed: {tok}")
    sys.exit("login timed out")


def token():
    tok = load()
    now = time.time()
    if tok and tok["expires_at"] - now > MIN_TTL:
        return tok
    if tok and tok.get("refresh_token") and tok["refresh_expires_at"] > now:
        new = post("token", grant_type="refresh_token", refresh_token=tok["refresh_token"])
        if "access_token" in new:
            return save(new)
    if not sys.stdin.isatty() and not sys.stderr.isatty():
        # headersHelper runs non-interactively: don't hang on a login nobody sees
        sys.exit("superset-mcp: not logged in, run mcp-token.py in a terminal")
    return device_login()


def main():
    arg = sys.argv[1] if len(sys.argv) > 1 else ""
    if arg == "--logout":
        if os.path.exists(CACHE):
            os.remove(CACHE)
        return
    access = token()["access_token"]
    print(json.dumps({"Authorization": f"Bearer {access}"}) if arg == "--header" else access)


if __name__ == "__main__":
    main()