"""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

    tokens = start_auth_flow()
    # tokens["refresh_token"] is stored in the Backend via /v1/miner/chutes-auth
"""

from __future__ import annotations

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

import httpx


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


# Default URLs
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 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

    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,
    )

    print("Chutes authentication successful!")
    return tokens
