"""OIDC Discovery and PKCE flow implementation.

This module handles Keycloak OIDC endpoint discovery and PKCE flow
for secure browser-based authentication, plus password grant (ROPC).
"""

from __future__ import annotations

import asyncio
import base64
import hashlib
import logging
import os
import secrets
import sys
from pathlib import Path
from typing import TYPE_CHECKING
from urllib.parse import urlencode

import httpx
from pydantic import BaseModel

from mcp_eregistrations_bpa.exceptions import AuthenticationError

if TYPE_CHECKING:
    from mcp_eregistrations_bpa.config import Config

logger = logging.getLogger(__name__)

# PKCE constants
PKCE_VERIFIER_LENGTH = 32  # bytes, produces 43 char base64url string


class OIDCConfig(BaseModel):
    """OIDC configuration discovered from well-known endpoint."""

    issuer: str
    authorization_endpoint: str
    token_endpoint: str
    userinfo_endpoint: str | None = None


async def discover_oidc_config(base_url: str) -> OIDCConfig:
    """Discover OIDC configuration from well-known endpoint.

    Args:
        base_url: The base URL for OIDC discovery. This is typically:
            - For Keycloak: https://login.example.org/realms/my-realm
            - Falls back to BPA URL if Keycloak URL not configured

    Returns:
        OIDCConfig with discovered endpoints.

    Raises:
        AuthenticationError: If discovery fails.
    """
    discovery_url = f"{base_url.rstrip('/')}/.well-known/openid-configuration"

    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(discovery_url, timeout=10.0)
            response.raise_for_status()
        except httpx.HTTPStatusError as e:
            raise AuthenticationError(
                f"Cannot discover Keycloak at {discovery_url}: "
                f"HTTP {e.response.status_code}. "
                "Verify KEYCLOAK_URL and KEYCLOAK_REALM are correct."
            ) from e
        except httpx.RequestError as e:
            raise AuthenticationError(
                f"Cannot connect to Keycloak at {discovery_url}: {e}. "
                "Verify network connectivity and that Keycloak is accessible."
            ) from e

        try:
            data = response.json()
            return OIDCConfig(
                issuer=data["issuer"],
                authorization_endpoint=data["authorization_endpoint"],
                token_endpoint=data["token_endpoint"],
                userinfo_endpoint=data.get("userinfo_endpoint"),
            )
        except (KeyError, ValueError) as e:
            raise AuthenticationError(
                f"Invalid OIDC configuration response: {e}. "
                "The URL may not be a valid Keycloak realm endpoint."
            ) from e


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

    Returns:
        Tuple of (code_verifier, code_challenge).
        The challenge is computed using S256 method.
    """
    # code_verifier: 43-128 characters, [A-Za-z0-9-._~]
    code_verifier = secrets.token_urlsafe(PKCE_VERIFIER_LENGTH)

    # code_challenge: SHA256(code_verifier) then base64url encode (no padding)
    digest = hashlib.sha256(code_verifier.encode()).digest()
    code_challenge = base64.urlsafe_b64encode(digest).decode().rstrip("=")

    return code_verifier, code_challenge


def generate_state() -> str:
    """Generate a cryptographically secure state parameter.

    Returns:
        A random state string for CSRF protection.
    """
    return secrets.token_urlsafe(16)


def build_authorization_url(
    authorization_endpoint: str,
    client_id: str,
    redirect_uri: str,
    code_challenge: str,
    state: str,
    scope: str = "openid email profile",
) -> str:
    """Build Keycloak authorization URL with PKCE.

    Args:
        authorization_endpoint: The Keycloak authorization endpoint.
        client_id: The OIDC client ID.
        redirect_uri: The local callback URL.
        code_challenge: The PKCE code challenge (S256).
        state: The state parameter for CSRF protection.
        scope: OAuth scopes to request.

    Returns:
        The complete authorization URL to open in browser.
    """
    params = {
        "response_type": "code",
        "client_id": client_id,
        "redirect_uri": redirect_uri,
        "scope": scope,
        "code_challenge": code_challenge,
        "code_challenge_method": "S256",
        "state": state,
    }
    return f"{authorization_endpoint}?{urlencode(params)}"


async def perform_browser_login(
    config: Config | None = None,
) -> dict[str, object]:
    """Perform browser-based OIDC login flow.

    This is the core login implementation that can be called from
    both the auth_login tool and auto-auth in ensure_authenticated.

    Args:
        config: Pre-resolved Config object. Falls back to load_config().

    Returns:
        dict: Authentication result with user email and session duration.

    Raises:
        AuthenticationError: If authentication fails.
    """
    import webbrowser

    from mcp_eregistrations_bpa.auth.callback import CallbackServer
    from mcp_eregistrations_bpa.auth.token_manager import exchange_code_for_tokens

    if config is None:
        from mcp_eregistrations_bpa.config import load_config

        config = load_config()

    logger.info("Starting browser login flow...")
    logger.debug("Config loaded: %s", config.oidc_discovery_url)

    # Discover OIDC endpoints
    logger.info("Discovering OIDC endpoints...")
    oidc_config = await discover_oidc_config(config.oidc_discovery_url)
    logger.debug("OIDC config: auth=%s", oidc_config.authorization_endpoint)

    # Generate PKCE pair and state
    code_verifier, code_challenge = generate_pkce_pair()
    state = generate_state()
    logger.debug("Generated PKCE pair and state")

    # Start callback server with instance info for the success page
    callback_server = CallbackServer(
        instance_info={
            "name": config.keycloak_realm or "BPA",
            "url": config.bpa_instance_url,
        }
    )
    callback_server.start()
    logger.info("Callback server started on port %d", callback_server.port)

    try:
        # Build authorization URL
        auth_url = build_authorization_url(
            authorization_endpoint=oidc_config.authorization_endpoint,
            client_id=config.keycloak_client_id,
            redirect_uri=callback_server.redirect_uri,
            code_challenge=code_challenge,
            state=state,
        )
        logger.debug("Auth URL built: %s...", auth_url[:80])

        # Open browser
        logger.info("Opening browser for authentication...")
        if not webbrowser.open(auth_url):
            logger.error("Failed to open browser")
            return {
                "error": True,
                "message": (
                    "Cannot open browser for authentication. "
                    f"Please open this URL manually: {auth_url}"
                ),
            }

        # Wait for callback
        logger.info("Waiting for OAuth callback...")
        code = await callback_server.wait_for_callback(expected_state=state)
        logger.info("Received authorization code")

        # Exchange code for tokens
        logger.info("Exchanging code for tokens...")
        token_response = await exchange_code_for_tokens(
            token_endpoint=oidc_config.token_endpoint,
            code=code,
            code_verifier=code_verifier,
            redirect_uri=callback_server.redirect_uri,
            client_id=config.keycloak_client_id,
        )
        logger.info("Token exchange successful")

        # Get token manager and store tokens
        from mcp_eregistrations_bpa.server import get_token_manager

        token_manager = get_token_manager(config.instance_id)
        token_manager.store_tokens(
            access_token=token_response.access_token,
            refresh_token=token_response.refresh_token,
            expires_in=token_response.expires_in,
            token_endpoint=oidc_config.token_endpoint,
            client_id=config.keycloak_client_id,
        )
        logger.info("Tokens stored for user: %s", token_manager.user_email)

        # Update the browser success page with user email
        callback_server.set_user_info(token_manager.user_email or "")
        # Brief delay so the browser JS can poll for user info
        await asyncio.sleep(1.5)

        result = {
            "success": True,
            "message": (
                f"Authenticated as {token_manager.user_email}. "
                f"Session valid for {token_manager.expires_in_minutes} minutes."
            ),
            "user_email": token_manager.user_email,
            "session_expires_in_minutes": token_manager.expires_in_minutes,
        }
        logger.info("Browser login complete, returning result")
        return result

    except AuthenticationError:
        logger.exception("Authentication error")
        raise
    except Exception as e:
        logger.exception("Unexpected error during authentication")
        raise AuthenticationError(
            f"Authentication failed: {e}. Please try again."
        ) from e
    finally:
        logger.debug("Stopping callback server")
        callback_server.stop()


def is_browser_available() -> bool:
    """Detect whether a browser can be opened in the current environment.

    Returns:
        True if browser-based login is likely to work.
    """
    # 1. Explicit headless flag (highest priority)
    if os.environ.get("MCP_HEADLESS", "").strip() == "1":
        return False

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

    # 3. Docker container detection
    if Path("/.dockerenv").exists():
        return False

    # 4. Dumb terminal
    if os.environ.get("TERM") == "dumb":
        return False

    # 5. macOS always has a browser (even via Terminal)
    if sys.platform == "darwin":
        return True

    # 6. Windows always has a browser
    if sys.platform == "win32":
        return True

    # 7. Linux: need a display server
    if os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY"):
        return True

    # 8. Default: no browser
    return False


async def perform_password_login(
    username: str,
    password: str,
    config: Config | None = None,
) -> dict[str, object]:
    """Authenticate via Keycloak Resource Owner Password Credentials grant.

    Args:
        username: BPA username (email).
        password: BPA password.
        config: Pre-resolved Config object. Falls back to load_config().

    Returns:
        dict with success, message, user_email, session_expires_in_minutes,
        auth_method.

    Raises:
        AuthenticationError: On invalid credentials or server errors.
    """
    logger.info("Attempting password grant for user: %s", username)

    if config is None:
        from mcp_eregistrations_bpa.config import load_config

        config = load_config()

    # Discover OIDC endpoints
    oidc_config = await discover_oidc_config(config.oidc_discovery_url)

    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(
                oidc_config.token_endpoint,
                data={
                    "grant_type": "password",
                    "client_id": config.keycloak_client_id,
                    "username": username,
                    "password": password,
                    "scope": "openid email profile",
                },
                headers={
                    "Content-Type": "application/x-www-form-urlencoded",
                },
                timeout=10.0,
            )
        except httpx.RequestError as e:
            raise AuthenticationError(
                f"Cannot reach authentication server: {e}. Check connectivity."
            ) from e

        if response.status_code == 401 or (
            response.status_code == 400 and "invalid_grant" in response.text
        ):
            raise AuthenticationError(
                "Invalid credentials. Verify your username and password. "
                "Do not retry automatically -- repeated failures may lock "
                "the account."
            )

        if response.status_code == 403:
            raise AuthenticationError(
                "Account disabled or not authorized for this client."
            )

        if response.status_code >= 400:
            detail = ""
            try:
                err = response.json()
                detail = err.get("error_description", err.get("error", ""))
            except Exception:
                pass
            raise AuthenticationError(
                f"Authentication failed (HTTP {response.status_code}): "
                f"{detail or 'Unknown error'}. Please try again later."
            )

        data = response.json()

    if "access_token" not in data:
        raise AuthenticationError(
            "Unexpected response from authentication server. Missing access_token."
        )

    # Store tokens via TokenManager
    from mcp_eregistrations_bpa.server import get_token_manager

    token_manager = get_token_manager(config.instance_id)
    token_manager.store_tokens(
        access_token=data["access_token"],
        refresh_token=data.get("refresh_token"),
        expires_in=data.get("expires_in", 300),
        token_endpoint=oidc_config.token_endpoint,
        client_id=config.keycloak_client_id,
    )

    logger.info("Password grant successful for user: %s", token_manager.user_email)

    return {
        "success": True,
        "message": (
            f"Authenticated as {token_manager.user_email}. "
            f"Session valid for {token_manager.expires_in_minutes} minutes."
        ),
        "user_email": token_manager.user_email,
        "session_expires_in_minutes": token_manager.expires_in_minutes,
        "auth_method": "password_grant",
    }
