#!/usr/bin/env python3
"""LeapAll OIDC login helper for LeapSkills Runtime Skills.

The helper opens the user's browser, completes Authorization Code + PKCE,
and stores an OIDC access token for the Skill's next runtime call.
"""

from __future__ import annotations

import base64
import hashlib
import json
import os
import secrets
import sys
import threading
import time
import webbrowser
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path
from typing import Any
from urllib.parse import parse_qs, urlencode, urlparse
from urllib.request import Request, urlopen

DEFAULT_ISSUER = "https://mercall.leapinfra.cn/api/v1/oidc"
DEFAULT_CLIENT_ID = "leapskills-runtime"
DEFAULT_REDIRECT_URI = "http://127.0.0.1:3345/callback"
DEFAULT_SCOPES = "openid profile profile.view offline_access"
CALLBACK_TIMEOUT_SECONDS = 300

def b64url(raw: bytes) -> str:
    return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")

def sha256_b64url(value: str) -> str:
    return b64url(hashlib.sha256(value.encode("utf-8")).digest())

def env_value(name: str, default: str = "") -> str:
    return os.environ.get(name, default).strip()

def default_token_path() -> Path:
    override = env_value("LEAPALL_OIDC_TOKEN_PATH")
    if override:
        return Path(override)
    if os.name == "nt":
        root = env_value("LOCALAPPDATA") or env_value("APPDATA") or str(Path.home())
        return Path(root) / "LeapSkills" / "leapall_oidc_access_token.txt"
    if sys.platform == "darwin":
        return Path.home() / "Library" / "Caches" / "LeapSkills" / "leapall_oidc_access_token.txt"
    root = env_value("XDG_CACHE_HOME")
    if root:
        return Path(root) / "leapskills" / "leapall_oidc_access_token.txt"
    return Path.home() / ".cache" / "leapskills" / "leapall_oidc_access_token.txt"

def token_path() -> Path:
    return default_token_path()

def fetch_json(url: str, method: str = "GET", headers: dict[str, str] | None = None, data: bytes | None = None) -> dict[str, Any]:
    request = Request(url=url, method=method, headers=headers or {}, data=data)
    with urlopen(request) as response:
        text = response.read().decode("utf-8")
        return json.loads(text) if text else {}

def wait_for_callback(redirect_uri: str, expected_state: str) -> str:
    parsed = urlparse(redirect_uri)
    host = parsed.hostname or "127.0.0.1"
    port = parsed.port or 80
    path = parsed.path or "/"
    event = threading.Event()
    result: dict[str, str] = {}

    class CallbackHandler(BaseHTTPRequestHandler):
        def do_GET(self) -> None:
            current = urlparse(self.path)
            if current.path != path:
                self.send_response(404)
                self.end_headers()
                return

            query = parse_qs(current.query)
            code = (query.get("code") or [""])[0]
            state = (query.get("state") or [""])[0]
            error = (query.get("error") or [""])[0]

            if error:
                result["error"] = error
                self.send_response(400)
                self.end_headers()
                self.wfile.write(f"Login error: {error}".encode("utf-8"))
                event.set()
                return

            if not code or state != expected_state:
                result["error"] = "invalid callback parameters"
                self.send_response(400)
                self.end_headers()
                self.wfile.write(b"Invalid callback parameters.")
                event.set()
                return

            result["code"] = code
            self.send_response(200)
            self.send_header("Content-Type", "text/html; charset=utf-8")
            self.end_headers()
            self.wfile.write("<html><body><h1>Login complete</h1><p>You can close this window.</p></body></html>".encode("utf-8"))
            event.set()

        def log_message(self, fmt: str, *args: Any) -> None:
            return

    server = HTTPServer((host, port), CallbackHandler)
    server.timeout = 1

    def serve() -> None:
        while not event.is_set():
            server.handle_request()

    thread = threading.Thread(target=serve, daemon=True)
    thread.start()
    deadline = time.time() + CALLBACK_TIMEOUT_SECONDS

    while time.time() < deadline:
        if event.wait(0.5):
            break

    event.set()
    server.server_close()
    thread.join(timeout=2)

    if "error" in result:
        raise RuntimeError(f"OIDC login failed: {result['error']}")
    if "code" not in result:
        raise RuntimeError("OIDC login timed out waiting for browser callback")
    return result["code"]

def main() -> None:
    print("Starting LeapAll OIDC login...", file=sys.stderr)
    issuer = env_value("LEAPALL_OIDC_ISSUER", DEFAULT_ISSUER).rstrip("/")
    client_id = env_value("LEAPALL_OIDC_CLIENT_ID", DEFAULT_CLIENT_ID)
    redirect_uri = env_value("LEAPALL_OIDC_REDIRECT_URI", DEFAULT_REDIRECT_URI)
    scopes = env_value("LEAPALL_OIDC_SCOPES", DEFAULT_SCOPES)

    print(f"OIDC client_id: {client_id}", file=sys.stderr)
    print(f"OIDC redirect_uri: {redirect_uri}", file=sys.stderr)

    discovery = fetch_json(f"{issuer}/.well-known/openid-configuration")
    authorization_endpoint = discovery["authorization_endpoint"]
    token_endpoint = discovery["token_endpoint"]

    state = secrets.token_urlsafe(24)
    code_verifier = secrets.token_urlsafe(48)[:64]
    code_challenge = sha256_b64url(code_verifier)

    authorization_url = authorization_endpoint + "?" + urlencode({
        "client_id": client_id,
        "redirect_uri": redirect_uri,
        "response_type": "code",
        "scope": scopes,
        "state": state,
        "code_challenge": code_challenge,
        "code_challenge_method": "S256",
    })

    print("Opening browser for LeapAll login...", file=sys.stderr)
    webbrowser.open(authorization_url)
    code = wait_for_callback(redirect_uri, state)

    body = urlencode({
        "grant_type": "authorization_code",
        "client_id": client_id,
        "redirect_uri": redirect_uri,
        "code": code,
        "code_verifier": code_verifier,
    }).encode("utf-8")

    token_response = fetch_json(
        token_endpoint,
        method="POST",
        headers={"Content-Type": "application/x-www-form-urlencoded"},
        data=body,
    )
    access_token = token_response["access_token"]

    path = token_path()
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(access_token + "\n", encoding="utf-8")
    try:
        path.chmod(0o600)
    except OSError:
        pass
    print(f"OIDC access token saved to {path}", file=sys.stderr)

if __name__ == "__main__":
    main()
