"""Chutes OAuth 2.0 + PKCE authentication flow for miner-funded inference.

Handles the browser-based OAuth flow so miners can authenticate with Chutes
and pass their refresh token to the Backend during agent submission.

Usage:
    from oro_sdk.chutes_auth import start_auth_flow, load_cached_tokens

    # Check for cached token first
    tokens = load_cached_tokens()
    if not tokens:
        tokens = start_auth_flow()

    # tokens["refresh_token"] is passed to the Backend on submit
"""

from __future__ import annotations

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

import httpx


class ChutesAuthError(Exception):
    """Raised when Chutes OAuth authentication fails."""


# Default paths and URLs
DEFAULT_TOKEN_CACHE = Path.home() / ".oro" / "chutes_tokens.json"
DEFAULT_REDIRECT_URI = "http://localhost:8765/callback"
DEFAULT_AUTH_URL = "https://api.chutes.ai/idp/authorize"
DEFAULT_TOKEN_URL = "https://api.chutes.ai/idp/token"

# Public OAuth client ID for the ORO SDK — same for all miners.
# PKCE replaces client_secret for public clients (RFC 7636).
DEFAULT_CLIENT_ID = os.environ.get("CHUTES_CLIENT_ID", "cid_qstkn2eb4uszxski1rdqk062")


def generate_pkce_pair() -> tuple[str, str]:
    """Generate a PKCE code_verifier and code_challenge pair.

    Returns:
        Tuple of (code_verifier, code_challenge) strings.
    """
    code_verifier = secrets.token_urlsafe(64)
    digest = hashlib.sha256(code_verifier.encode()).digest()
    code_challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode()
    return code_verifier, code_challenge


def _build_authorize_url(
    client_id: str,
    redirect_uri: str,
    code_challenge: str,
    state: str,
    authorize_url: str = DEFAULT_AUTH_URL,
    scope: str = "chutes:invoke",
) -> str:
    """Build the Chutes authorization URL with PKCE parameters."""
    params = {
        "response_type": "code",
        "client_id": client_id,
        "redirect_uri": redirect_uri,
        "code_challenge": code_challenge,
        "code_challenge_method": "S256",
        "state": state,
        "scope": scope,
    }
    return f"{authorize_url}?{urlencode(params)}"


def exchange_code(
    code: str,
    verifier: str,
    client_id: str,
    redirect_uri: str,
    token_url: str = DEFAULT_TOKEN_URL,
) -> dict[str, Any]:
    """Exchange an authorization code for access and refresh tokens.

    Uses PKCE (code_verifier) instead of client_secret per RFC 7636.
    """
    with httpx.Client(timeout=15.0) as client:
        response = client.post(
            token_url,
            data={
                "grant_type": "authorization_code",
                "code": code,
                "code_verifier": verifier,
                "client_id": client_id,
                "redirect_uri": redirect_uri,
            },
        )
        response.raise_for_status()
        return response.json()



def load_cached_tokens(path: Path = DEFAULT_TOKEN_CACHE) -> dict[str, Any] | None:
    """Load cached Chutes tokens from disk.

    Returns:
        Token dict if cache exists and is valid, None otherwise.
    """
    if not path.exists():
        return None
    try:
        data = json.loads(path.read_text())
        if "refresh_token" in data:
            return data
    except (json.JSONDecodeError, OSError):
        pass
    return None


def save_cached_tokens(tokens: dict[str, Any], path: Path = DEFAULT_TOKEN_CACHE) -> None:
    """Save tokens to the cache file with restricted permissions.

    Creates ~/.oro/ (0o700) if needed. File is created atomically with 0o600.
    """
    path.parent.mkdir(parents=True, exist_ok=True)
    os.chmod(str(path.parent), 0o700)
    fd = os.open(str(path), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    try:
        os.write(fd, json.dumps(tokens, indent=2).encode())
    finally:
        os.close(fd)


def start_auth_flow(
    client_id: str = DEFAULT_CLIENT_ID,
    redirect_uri: str = DEFAULT_REDIRECT_URI,
    authorize_url: str = DEFAULT_AUTH_URL,
    token_url: str = DEFAULT_TOKEN_URL,
) -> dict[str, Any]:
    """Run the full OAuth 2.0 + PKCE flow interactively.

    1. Generates PKCE pair and CSRF state token
    2. Opens browser to Chutes authorization page
    3. Starts a local HTTP server to receive the callback
    4. Exchanges the authorization code for tokens
    5. Caches the tokens to ~/.oro/chutes_tokens.json

    Args:
        client_id: OAuth client ID (default: ORO SDK public client).
        redirect_uri: Localhost callback URI.
        authorize_url: Chutes authorization endpoint.
        token_url: Chutes token endpoint.

    Returns:
        Dict with "access_token", "refresh_token", etc.
    """
    verifier, challenge = generate_pkce_pair()
    state = secrets.token_urlsafe(32)
    auth_url = _build_authorize_url(client_id, redirect_uri, challenge, state, authorize_url)

    parsed = urlparse(redirect_uri)
    port = parsed.port or 8765

    auth_code: dict[str, str | None] = {"code": None, "error": None}

    class CallbackHandler(BaseHTTPRequestHandler):
        def do_GET(self) -> None:
            query = parse_qs(urlparse(self.path).query)

            # Verify state to prevent CSRF
            received_state = query.get("state", [None])[0]
            if received_state != state:
                auth_code["error"] = "State mismatch — possible CSRF attack"
                self.send_response(400)
                self.send_header("Content-Type", "text/html")
                self.end_headers()
                self.wfile.write(b"<html><body><h2>State mismatch</h2></body></html>")
                threading.Thread(target=server.shutdown).start()
                return

            if "code" in query:
                auth_code["code"] = query["code"][0]
                self.send_response(200)
                self.send_header("Content-Type", "text/html")
                self.end_headers()
                self.wfile.write(
                    b"<html><body><h2>Authorization successful!</h2>"
                    b"<p>You can close this tab and return to the terminal.</p>"
                    b"</body></html>"
                )
            elif "error" in query:
                auth_code["error"] = query.get("error_description", query["error"])[0]
                self.send_response(400)
                self.send_header("Content-Type", "text/html")
                self.end_headers()
                self.wfile.write(
                    f"<html><body><h2>Authorization failed</h2>"
                    f"<p>{auth_code['error']}</p></body></html>".encode()
                )
            else:
                # Ignore favicon, preflight, etc.
                self.send_response(204)
                self.end_headers()
                return

            threading.Thread(target=server.shutdown).start()

        def log_message(self, format: str, *args: Any) -> None:
            pass  # Suppress HTTP server logs

    server = HTTPServer(("localhost", port), CallbackHandler)

    print("Opening browser for Chutes authentication...")
    print(f"If the browser doesn't open, visit: {auth_url}")
    webbrowser.open(auth_url)

    server_thread = threading.Thread(target=server.serve_forever)
    server_thread.start()
    server_thread.join(timeout=300)
    server.shutdown()

    if auth_code["error"]:
        raise ChutesAuthError(f"Authentication failed: {auth_code['error']}")

    if not auth_code["code"]:
        raise ChutesAuthError("Authentication timed out. Please try again.")

    print("Exchanging authorization code for tokens...")
    tokens = exchange_code(
        code=auth_code["code"],
        verifier=verifier,
        client_id=client_id,
        redirect_uri=redirect_uri,
        token_url=token_url,
    )

    save_cached_tokens(tokens)
    print("Chutes authentication successful! Token cached to ~/.oro/chutes_tokens.json")
    return tokens
