#!/usr/bin/env python3
"""
OpenAI OAuth dynamic provider for maki (device code flow).

Install: cp scripts/providers/openai-oauth ~/.maki/providers/openai-oauth && chmod +x ~/.maki/providers/openai-oauth
Usage:   maki auth login openai-oauth
         maki -m openai-oauth/gpt-4.1

Requires Python 3.8+ with no external dependencies.
Stores tokens at ~/.maki/auth/openai.json
"""

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

CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
DEVICE_CODE_URL = "https://auth.openai.com/api/accounts/deviceauth/usercode"
DEVICE_TOKEN_URL = "https://auth.openai.com/api/accounts/deviceauth/token"
OAUTH_TOKEN_URL = "https://auth.openai.com/oauth/token"
DEVICE_AUTH_URL = "https://auth.openai.com/codex/device"
REDIRECT_URI = "https://auth.openai.com/deviceauth/callback"
REFRESH_BUFFER_SECS = 60
POLL_TIMEOUT = 300

TOKEN_FILE = os.path.expanduser("~/.maki/auth/openai.json")


def load_tokens():
    try:
        with open(TOKEN_FILE) as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        return None


def save_tokens(tokens):
    os.makedirs(os.path.dirname(TOKEN_FILE), exist_ok=True)
    tmp = TOKEN_FILE + ".tmp"
    with open(tmp, "w") as f:
        json.dump(tokens, f, indent=2)
    os.chmod(tmp, 0o600)
    os.replace(tmp, TOKEN_FILE)


def now_ms():
    return int(time.time() * 1000)


def is_expired(tokens):
    return now_ms() + REFRESH_BUFFER_SECS * 1000 >= tokens["expires"]


def post_json(url, data):
    body = json.dumps(data).encode()
    req = urllib.request.Request(
        url,
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        return json.loads(resp.read())


def post_form(url, data):
    body = urllib.parse.urlencode(data).encode()
    req = urllib.request.Request(
        url,
        data=body,
        headers={"Content-Type": "application/x-www-form-urlencoded"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        return json.loads(resp.read())


def build_resolved(tokens):
    return {
        "headers": {
            "authorization": f"Bearer {tokens['access']}",
        },
    }


def do_refresh(tokens):
    resp = post_form(OAUTH_TOKEN_URL, {
        "grant_type": "refresh_token",
        "refresh_token": tokens["refresh"],
        "client_id": CLIENT_ID,
    })
    new_tokens = {
        "access": resp["access_token"],
        "refresh": resp["refresh_token"],
        "expires": now_ms() + resp.get("expires_in", 3600) * 1000,
    }
    save_tokens(new_tokens)
    return new_tokens


def cmd_info():
    print(json.dumps({
        "display_name": "OpenAI OAuth",
        "base": "openai",
        "has_auth": True,
    }))


def cmd_resolve():
    tokens = load_tokens()
    if not tokens:
        print("not authenticated, run: maki auth login openai-oauth", file=sys.stderr)
        sys.exit(1)
    if is_expired(tokens):
        try:
            tokens = do_refresh(tokens)
        except Exception as e:
            print(f"token refresh failed: {e}", file=sys.stderr)
            sys.exit(1)
    print(json.dumps(build_resolved(tokens)))


def cmd_refresh():
    tokens = load_tokens()
    if not tokens:
        print("no tokens to refresh", file=sys.stderr)
        sys.exit(1)
    try:
        tokens = do_refresh(tokens)
    except Exception as e:
        print(f"token refresh failed: {e}", file=sys.stderr)
        sys.exit(1)
    print(json.dumps(build_resolved(tokens)))


def cmd_login():
    resp = post_json(DEVICE_CODE_URL, {"client_id": CLIENT_ID})
    device_auth_id = resp["device_auth_id"]
    user_code = resp["user_code"]
    interval = max(int(resp.get("interval", "5")), 1) + 3

    print(f"Open this URL in your browser:\n\n  {DEVICE_AUTH_URL}\n")
    print(f"Enter code: {user_code}\n")
    print("Waiting for authorization...")

    deadline = time.time() + POLL_TIMEOUT
    while time.time() < deadline:
        time.sleep(interval)
        try:
            poll_resp = post_json(DEVICE_TOKEN_URL, {
                "device_auth_id": device_auth_id,
                "user_code": user_code,
            })
            break
        except urllib.error.HTTPError as e:
            if e.code in (403, 404):
                continue
            raise
    else:
        print("device authorization timed out", file=sys.stderr)
        sys.exit(1)

    token_resp = post_form(OAUTH_TOKEN_URL, {
        "grant_type": "authorization_code",
        "code": poll_resp["authorization_code"],
        "redirect_uri": REDIRECT_URI,
        "client_id": CLIENT_ID,
        "code_verifier": poll_resp["code_verifier"],
    })

    tokens = {
        "access": token_resp["access_token"],
        "refresh": token_resp["refresh_token"],
        "expires": now_ms() + token_resp.get("expires_in", 3600) * 1000,
    }
    save_tokens(tokens)
    print("Authenticated successfully.")


def cmd_logout():
    try:
        os.remove(TOKEN_FILE)
        print("Logged out of OpenAI.")
    except FileNotFoundError:
        print("Not currently logged in to OpenAI.")


def main():
    if len(sys.argv) < 2:
        print(f"usage: {sys.argv[0]} <info|resolve|refresh|login|logout>", file=sys.stderr)
        sys.exit(1)

    cmd = sys.argv[1]
    commands = {
        "info": cmd_info,
        "resolve": cmd_resolve,
        "refresh": cmd_refresh,
        "login": cmd_login,
        "logout": cmd_logout,
    }

    fn = commands.get(cmd)
    if not fn:
        print(f"unknown command: {cmd}", file=sys.stderr)
        sys.exit(1)
    fn()


if __name__ == "__main__":
    main()
