#!/usr/bin/env python3
"""OpenAI OAuth token management with self-contained login flows.

OAuth flow ported from OpenClaw (MIT, github.com/openclaw/openclaw).

Supports three authentication methods (checked in order):
1. Own OAuth token stored in ~/.second-opinion-config under openai_oauth key
2. Codex CLI OAuth token from ~/.codex/auth.json (fallback)
3. Standard API key from keyring (pay-as-you-go, handled in server.py)

Two OAuth flows are supported, auto-detected by environment:
- Browser PKCE flow (desktop with GUI): opens browser, local callback server
- Device code flow (headless/SSH): user enters code on another device
"""
import base64
import hashlib
import html as html_module
import json
import logging
import os
import secrets
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from datetime import datetime, timezone
from http.server import HTTPServer, BaseHTTPRequestHandler
from pathlib import Path

logger = logging.getLogger("second-opinion")

# OAuth constants
CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
AUTHORIZE_URL = "https://auth.openai.com/oauth/authorize"
TOKEN_URL = "https://auth.openai.com/oauth/token"
REDIRECT_URI = "http://localhost:1455/auth/callback"
SCOPE = "openid profile email offline_access"
DEVICE_AUTH_URL = "https://auth.openai.com/api/accounts/deviceauth/usercode"
DEVICE_TOKEN_URL = "https://auth.openai.com/api/accounts/deviceauth/token"
DEVICE_VERIFY_URL = "https://auth.openai.com/codex/device"

# Legacy Codex CLI support
CODEX_AUTH_FILE = Path.home() / ".codex" / "auth.json"


def _safe_error_from_body(body_text: str) -> str:
    """Extract only safe OAuth error fields from response body."""
    try:
        data = json.loads(body_text)
        error = data.get("error", "unknown")
        desc = data.get("error_description", "")
        return f"{error}: {desc}" if desc else str(error)
    except (json.JSONDecodeError, AttributeError):
        return "non-JSON response"


# ---------------------------------------------------------------------------
# PKCE helpers
# ---------------------------------------------------------------------------

def _generate_pkce() -> tuple[str, str]:
    """Generate PKCE verifier and challenge.

    Returns:
        (verifier, challenge) tuple where verifier is 43-char base64url string
        and challenge is SHA-256 of verifier, base64url-encoded.
    """
    verifier_bytes = secrets.token_bytes(32)
    verifier = base64.urlsafe_b64encode(verifier_bytes).rstrip(b"=").decode("ascii")
    challenge_digest = hashlib.sha256(verifier.encode("ascii")).digest()
    challenge = base64.urlsafe_b64encode(challenge_digest).rstrip(b"=").decode("ascii")
    return verifier, challenge


# ---------------------------------------------------------------------------
# Token exchange (shared by both flows)
# ---------------------------------------------------------------------------

def _exchange_code_for_tokens(code: str, verifier: str) -> dict | None:
    """Exchange authorization code for tokens via TOKEN_URL.

    Args:
        code: Authorization code from callback or device auth.
        verifier: PKCE code verifier.

    Returns:
        Dict with access_token, refresh_token, id_token, or None on failure.
    """
    body = urllib.parse.urlencode({
        "grant_type": "authorization_code",
        "client_id": CLIENT_ID,
        "code": code,
        "code_verifier": verifier,
        "redirect_uri": REDIRECT_URI,
    }).encode("utf-8")

    req = urllib.request.Request(
        TOKEN_URL,
        data=body,
        headers={"Content-Type": "application/x-www-form-urlencoded"},
        method="POST",
    )

    try:
        with urllib.request.urlopen(req, timeout=15) as resp:
            data = json.loads(resp.read().decode("utf-8"))
        if not data.get("access_token"):
            logger.error("Token exchange response missing access_token")
            return None
        return data
    except urllib.error.HTTPError as e:
        body_text = e.read().decode("utf-8", errors="replace") if e.fp else ""
        logger.error("Token exchange failed (HTTP %d): %s", e.code, _safe_error_from_body(body_text))
        return None
    except (urllib.error.URLError, OSError, json.JSONDecodeError) as e:
        logger.error("Token exchange failed: %s", e)
        return None


# ---------------------------------------------------------------------------
# Browser PKCE flow (Flow A)
# ---------------------------------------------------------------------------

class _CallbackHandler(BaseHTTPRequestHandler):
    """HTTP handler for OAuth callback on localhost:1455."""

    authorization_code: str | None = None
    received_state: str | None = None

    def do_GET(self):
        parsed = urllib.parse.urlparse(self.path)
        if parsed.path != "/auth/callback":
            self.send_response(404)
            self.end_headers()
            return

        params = urllib.parse.parse_qs(parsed.query)

        # Check for OAuth error response first
        error = params.get("error", [None])[0]
        if error:
            error_desc = params.get("error_description", ["Unknown error"])[0]
            safe_desc = html_module.escape(error_desc)
            _CallbackHandler.authorization_code = None
            _CallbackHandler.received_state = params.get("state", [None])[0]
            self.send_response(200)
            self.send_header("Content-Type", "text/html")
            self.end_headers()
            error_html = (
                "<html><body style='font-family:sans-serif;text-align:center;padding:60px'>"
                "<h2>Login failed</h2>"
                f"<p>{html_module.escape(error)}: {safe_desc}</p>"
                "<p>You can close this tab and try again in the terminal.</p>"
                "</body></html>"
            )
            self.wfile.write(error_html.encode("utf-8"))
            return

        _CallbackHandler.received_state = params.get("state", [None])[0]
        _CallbackHandler.authorization_code = params.get("code", [None])[0]

        self.send_response(200)
        self.send_header("Content-Type", "text/html")
        self.end_headers()
        success_html = (
            "<html><body style='font-family:sans-serif;text-align:center;padding:60px'>"
            "<h2>Login successful!</h2>"
            "<p>You can close this tab and return to the terminal.</p>"
            "</body></html>"
        )
        self.wfile.write(success_html.encode("utf-8"))

    def log_message(self, format, *args):
        # Suppress default HTTP logging
        pass


def _start_callback_server(expected_state: str, timeout: float = 120) -> str | None:
    """Start local HTTP server and wait for OAuth callback.

    Args:
        expected_state: The state parameter to validate against.
        timeout: Seconds to wait before giving up.

    Returns:
        Authorization code string, or None on timeout/state mismatch.
    """
    _CallbackHandler.authorization_code = None
    _CallbackHandler.received_state = None

    try:
        server = HTTPServer(("127.0.0.1", 1455), _CallbackHandler)
    except OSError as e:
        logger.error("Could not start callback server on port 1455: %s", e)
        return None

    server.timeout = timeout

    # Handle one request (the callback)
    server.handle_request()
    server.server_close()

    if _CallbackHandler.received_state != expected_state:
        logger.error("OAuth state mismatch")
        return None

    return _CallbackHandler.authorization_code


def _browser_login() -> dict | None:
    """Run the browser-based PKCE OAuth flow.

    Opens browser to OpenAI auth, starts local callback server,
    exchanges code for tokens.

    Returns:
        Token dict with access_token/refresh_token/id_token, or None on failure.
    """
    import webbrowser

    verifier, challenge = _generate_pkce()
    state = secrets.token_hex(16)

    params = urllib.parse.urlencode({
        "response_type": "code",
        "client_id": CLIENT_ID,
        "redirect_uri": REDIRECT_URI,
        "scope": SCOPE,
        "code_challenge": challenge,
        "code_challenge_method": "S256",
        "state": state,
    })
    auth_url = f"{AUTHORIZE_URL}?{params}"

    print("  Opening browser for OpenAI login...")
    webbrowser.open(auth_url)
    print("  Waiting for authorization...")

    code = _start_callback_server(state, timeout=120)
    if not code:
        logger.error("Browser login: no authorization code received")
        return None

    return _exchange_code_for_tokens(code, verifier)


# ---------------------------------------------------------------------------
# Device code flow (Flow B)
# ---------------------------------------------------------------------------

def _request_device_code() -> dict | None:
    """Request a device code from OpenAI.

    Returns:
        Dict with device_auth_id, user_code, interval, or None on failure.
    """
    body = json.dumps({"client_id": CLIENT_ID}).encode("utf-8")
    req = urllib.request.Request(
        DEVICE_AUTH_URL,
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )

    try:
        with urllib.request.urlopen(req, timeout=15) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        body_text = e.read().decode("utf-8", errors="replace") if e.fp else ""
        logger.error("Device code request failed (HTTP %d): %s", e.code, _safe_error_from_body(body_text))
        return None
    except (urllib.error.URLError, OSError, json.JSONDecodeError) as e:
        logger.error("Device code request failed: %s", e)
        return None


def _poll_device_auth(
    device_auth_id: str,
    user_code: str,
    interval: int = 5,
    timeout: float = 900,
) -> dict | None:
    """Poll for device authorization completion.

    Args:
        device_auth_id: Device auth session ID.
        user_code: The code the user enters.
        interval: Seconds between polls.
        timeout: Maximum seconds to wait.

    Returns:
        Dict with authorization_code and code_verifier, or None on timeout.
    """
    body = json.dumps({
        "device_auth_id": device_auth_id,
        "user_code": user_code,
    }).encode("utf-8")

    deadline = time.monotonic() + timeout

    while time.monotonic() < deadline:
        req = urllib.request.Request(
            DEVICE_TOKEN_URL,
            data=body,
            headers={"Content-Type": "application/json"},
            method="POST",
        )

        try:
            with urllib.request.urlopen(req, timeout=15) as resp:
                data = json.loads(resp.read().decode("utf-8"))
                return data  # Success: has authorization_code + code_verifier
        except urllib.error.HTTPError as e:
            body_text = e.read().decode("utf-8", errors="replace") if e.fp else ""
            try:
                err_data = json.loads(body_text)
                err_code = err_data.get("error", "")
            except (json.JSONDecodeError, AttributeError):
                err_code = ""

            if err_code in ("authorization_pending", "") and e.code in (403, 404):
                time.sleep(interval)
                continue
            elif err_code == "slow_down":
                interval = min(interval + 5, 30)
                time.sleep(interval)
                continue
            elif err_code in ("access_denied", "expired_token"):
                logger.error("Device auth %s", err_code)
                return None
            else:
                logger.error("Device auth poll failed (HTTP %d): %s", e.code, _safe_error_from_body(body_text))
                return None
        except (urllib.error.URLError, OSError, json.JSONDecodeError) as e:
            logger.error("Device auth poll error: %s", e)
            return None

    logger.error("Device auth timed out after %.0f seconds", timeout)
    return None


def _device_code_login() -> dict | None:
    """Run the device code OAuth flow (headless environments).

    Requests a device code, prompts user to authorize on another device,
    polls for completion, then exchanges code for tokens.

    Returns:
        Token dict with access_token/refresh_token/id_token, or None on failure.
    """
    device_data = _request_device_code()
    if not device_data:
        return None

    device_auth_id = device_data.get("device_auth_id")
    user_code = device_data.get("user_code")
    interval = device_data.get("interval", 5)

    if not device_auth_id or not user_code:
        logger.error("Device code response missing required fields")
        return None

    print(f"\n  Visit {DEVICE_VERIFY_URL} and enter code: {user_code}\n")
    print("  Waiting for authorization...")

    poll_result = _poll_device_auth(device_auth_id, user_code, interval)
    if not poll_result:
        return None

    auth_code = poll_result.get("authorization_code")
    code_verifier = poll_result.get("code_verifier")

    if not auth_code or not code_verifier:
        logger.error("Device auth response missing authorization_code or code_verifier")
        return None

    return _exchange_code_for_tokens(auth_code, code_verifier)


# ---------------------------------------------------------------------------
# Manual URL paste fallback
# ---------------------------------------------------------------------------

def _manual_url_login() -> dict | None:
    """Fallback when device code is blocked: show browser URL, ask user to paste redirect.

    Returns:
        Token dict or None on failure.
    """
    verifier, challenge = _generate_pkce()
    state = secrets.token_hex(16)

    params = urllib.parse.urlencode({
        "response_type": "code",
        "client_id": CLIENT_ID,
        "redirect_uri": REDIRECT_URI,
        "scope": SCOPE,
        "code_challenge": challenge,
        "code_challenge_method": "S256",
        "state": state,
    })
    auth_url = f"{AUTHORIZE_URL}?{params}"

    print(f"\n  Open this URL in a browser:\n  {auth_url}\n")
    print("  After signing in, paste the full redirect URL here")
    print("  (it will start with http://localhost:1455/auth/callback?...)")

    try:
        redirect_url = input("  URL: ").strip()
    except (EOFError, KeyboardInterrupt):
        return None

    if not redirect_url:
        return None

    parsed = urllib.parse.urlparse(redirect_url)

    # Validate redirect URI
    if parsed.scheme != "http" or parsed.netloc not in ("localhost:1455", "127.0.0.1:1455") or parsed.path != "/auth/callback":
        logger.error("Invalid redirect URL: expected http://localhost:1455/auth/callback")
        print("  Error: URL must be the OAuth callback (http://localhost:1455/auth/callback?...)")
        return None

    params_dict = urllib.parse.parse_qs(parsed.query)

    # Check for OAuth error in redirect
    error = params_dict.get("error", [None])[0]
    if error:
        error_desc = params_dict.get("error_description", [""])[0]
        logger.error("OAuth error in redirect: %s - %s", error, error_desc)
        print(f"  Error: authorization failed ({error}: {error_desc})")
        return None

    received_state = params_dict.get("state", [None])[0]
    code = params_dict.get("code", [None])[0]

    if received_state != state:
        logger.error("OAuth state mismatch in manual URL")
        print("  Error: state mismatch. Please try again.")
        return None

    if not code:
        logger.error("No authorization code in pasted URL")
        print("  Error: no authorization code found in URL.")
        return None

    return _exchange_code_for_tokens(code, verifier)


# ---------------------------------------------------------------------------
# Environment detection
# ---------------------------------------------------------------------------

def _is_headless() -> bool:
    """Detect if running in a headless environment (no browser available).

    Returns True for SSH sessions, Docker, WSL without display, etc.
    On Windows, always returns False (always has browser).
    """
    if sys.platform == "win32":
        return False

    # SSH session
    if os.environ.get("SSH_CONNECTION") or os.environ.get("SSH_TTY"):
        return True

    if sys.platform == "darwin":
        # macOS: check for WindowServer
        if not os.environ.get("__CFBundleIdentifier"):
            try:
                import subprocess
                result = subprocess.run(
                    ["pgrep", "-x", "WindowServer"],
                    capture_output=True, timeout=2,
                )
                if result.returncode != 0:
                    return True
            except Exception:
                logger.debug("WindowServer check failed", exc_info=True)
        return False

    # Linux: check DISPLAY or WAYLAND_DISPLAY
    if not os.environ.get("DISPLAY") and not os.environ.get("WAYLAND_DISPLAY"):
        return True

    return False


# ---------------------------------------------------------------------------
# OAuth login orchestrator
# ---------------------------------------------------------------------------

def oauth_login() -> bool:
    """Run the appropriate OAuth flow based on environment.

    Auto-detects desktop vs headless and picks the right flow.
    If device code is blocked (HTTP 400), falls back to manual URL paste.

    Returns:
        True if login succeeded and tokens were stored, False otherwise.
    """
    from second_opinion_mcp.credentials import CredentialManager

    if _is_headless():
        tokens = _device_code_login()
        if tokens is None:
            # Device code might be blocked, try manual URL paste
            print("  Device code flow unavailable. Trying manual URL paste...")
            tokens = _manual_url_login()
    else:
        tokens = _browser_login()

    if not tokens or not tokens.get("access_token"):
        return False

    # Store tokens in config file
    oauth_data = {
        "access_token": tokens["access_token"],
        "refresh_token": tokens.get("refresh_token", ""),
        "id_token": tokens.get("id_token", ""),
        "obtained_at": datetime.now(timezone.utc).isoformat(),
    }

    manager = CredentialManager()
    manager.set_config_section("openai_oauth", oauth_data)
    print("  Login successful!")
    return True


# ---------------------------------------------------------------------------
# Token refresh (form-encoded, used by both own OAuth and Codex)
# ---------------------------------------------------------------------------

def refresh_access_token(refresh_token: str) -> dict | None:
    """Refresh the OAuth access token using the refresh token.

    Uses application/x-www-form-urlencoded as required by the OAuth spec.

    Args:
        refresh_token: The OAuth refresh token.

    Returns:
        Dict with new tokens, or None on failure.
    """
    body = urllib.parse.urlencode({
        "grant_type": "refresh_token",
        "refresh_token": refresh_token,
        "client_id": CLIENT_ID,
    }).encode("utf-8")

    req = urllib.request.Request(
        TOKEN_URL,
        data=body,
        headers={"Content-Type": "application/x-www-form-urlencoded"},
        method="POST",
    )

    try:
        with urllib.request.urlopen(req, timeout=10) as resp:
            data = json.loads(resp.read().decode("utf-8"))

        access_token = data.get("access_token")
        if not access_token:
            logger.error("Token refresh response missing access_token")
            return None

        new_refresh = data.get("refresh_token", refresh_token)
        logger.info("Successfully refreshed OpenAI access token")
        return {
            "access_token": access_token,
            "refresh_token": new_refresh,
            "id_token": data.get("id_token"),
        }
    except urllib.error.HTTPError as e:
        body_text = e.read().decode("utf-8", errors="replace") if e.fp else ""
        logger.error("Token refresh failed (HTTP %d): %s", e.code, _safe_error_from_body(body_text))
        return None
    except (urllib.error.URLError, OSError, json.JSONDecodeError) as e:
        logger.error("Token refresh failed: %s", e)
        return None


# ---------------------------------------------------------------------------
# JWT helpers
# ---------------------------------------------------------------------------

def _decode_jwt_exp(token: str) -> float | None:
    """Extract expiration timestamp from a JWT access token."""
    try:
        payload_b64 = token.split(".")[1]
        payload_b64 += "=" * (4 - len(payload_b64) % 4)
        payload = json.loads(base64.urlsafe_b64decode(payload_b64))
        return float(payload["exp"])
    except Exception:
        return None


def _is_token_expired(token: str) -> bool:
    """Check if a JWT access token is expired (with 60s buffer)."""
    exp = _decode_jwt_exp(token)
    if exp is None:
        return True
    return time.time() >= (exp - 60)


# ---------------------------------------------------------------------------
# Own OAuth token management
# ---------------------------------------------------------------------------

def get_own_oauth_token() -> str | None:
    """Get a valid OAuth token from own stored tokens.

    Reads from ~/.second-opinion-config openai_oauth section.
    Auto-refreshes if expired.

    Returns:
        Valid access token string, or None if unavailable.
    """
    from second_opinion_mcp.credentials import CredentialManager

    manager = CredentialManager()
    oauth_data = manager.get_config_section("openai_oauth")
    if not oauth_data:
        return None

    access_token = oauth_data.get("access_token")
    refresh_token = oauth_data.get("refresh_token")

    if not access_token and not refresh_token:
        return None

    # If token is still valid, return it
    if access_token and not _is_token_expired(access_token):
        logger.debug("Using stored OAuth access token")
        return access_token

    # Token expired, try refresh
    if refresh_token:
        logger.info("Own OAuth token expired, attempting refresh")
        result = refresh_access_token(refresh_token)
        if result:
            # Update stored tokens
            oauth_data["access_token"] = result["access_token"]
            if result.get("refresh_token"):
                oauth_data["refresh_token"] = result["refresh_token"]
            if result.get("id_token"):
                oauth_data["id_token"] = result["id_token"]
            oauth_data["obtained_at"] = datetime.now(timezone.utc).isoformat()
            manager.set_config_section("openai_oauth", oauth_data)
            return result["access_token"]

    logger.warning("Could not obtain valid own OAuth token")
    return None


# ---------------------------------------------------------------------------
# Codex CLI fallback (kept for backwards compatibility)
# ---------------------------------------------------------------------------

def read_codex_auth() -> dict | None:
    """Read and parse ~/.codex/auth.json.

    Returns:
        Parsed auth data dict, or None if file doesn't exist or is invalid.
    """
    if not CODEX_AUTH_FILE.exists():
        logger.debug("Codex auth file not found: %s", CODEX_AUTH_FILE)
        return None

    try:
        data = json.loads(CODEX_AUTH_FILE.read_text(encoding="utf-8"))
        if not isinstance(data, dict):
            logger.warning("Codex auth file has unexpected format")
            return None
        return data
    except (json.JSONDecodeError, OSError) as e:
        logger.warning("Failed to read Codex auth file: %s", e)
        return None


def _update_codex_auth(access_token: str, refresh_token: str, id_token: str | None = None) -> None:
    """Update ~/.codex/auth.json with new token data, preserving existing structure."""
    auth_data = read_codex_auth() or {}

    tokens = auth_data.get("tokens", {})
    tokens["access_token"] = access_token
    tokens["refresh_token"] = refresh_token
    if id_token:
        tokens["id_token"] = id_token
    auth_data["tokens"] = tokens
    auth_data["last_refresh"] = datetime.now(timezone.utc).isoformat()

    try:
        CODEX_AUTH_FILE.parent.mkdir(parents=True, exist_ok=True)
        CODEX_AUTH_FILE.write_text(
            json.dumps(auth_data, indent=2),
            encoding="utf-8",
        )
        logger.debug("Updated Codex auth file")
    except OSError as e:
        logger.warning("Failed to update Codex auth file: %s", e)


def _get_codex_token() -> str | None:
    """Get a valid token from Codex CLI auth (fallback).

    Returns:
        Valid access token string, or None if unavailable.
    """
    auth_data = read_codex_auth()
    if not auth_data:
        return None

    tokens = auth_data.get("tokens", {})
    access_token = tokens.get("access_token")
    refresh_token = tokens.get("refresh_token")

    if not access_token and not refresh_token:
        logger.debug("No tokens found in Codex auth file")
        return None

    if access_token and not _is_token_expired(access_token):
        logger.debug("Using existing Codex access token")
        return access_token

    if refresh_token:
        logger.info("Codex access token expired, attempting refresh")
        result = refresh_access_token(refresh_token)
        if result:
            _update_codex_auth(
                result["access_token"],
                result["refresh_token"],
                result.get("id_token"),
            )
            return result["access_token"]

    logger.warning("Could not obtain valid token from Codex CLI")
    return None


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------

def get_openai_token() -> str | None:
    """Get a valid OpenAI access token.

    Priority:
    1. Own OAuth token (from ~/.second-opinion-config openai_oauth)
    2. Codex CLI token (from ~/.codex/auth.json)

    Returns:
        Valid access token string, or None if unavailable.
    """
    # Try own OAuth first
    token = get_own_oauth_token()
    if token:
        return token

    # Fall back to Codex CLI
    token = _get_codex_token()
    if token:
        logger.info("Using OpenAI token from Codex CLI auth")
        return token

    return None
