"""Auth routes — Google OAuth + SSO login flows for dashboard.

- GET /auth/login → email-first login routing (SSO or Google OAuth)
- GET /auth/callback → Google OAuth code exchange + session creation
- GET /auth/sso/callback → SSO code exchange + session creation
- POST /auth/logout → CSRF-validated session teardown
"""

from __future__ import annotations

import base64
import json
import logging
import secrets
import time
from datetime import UTC, datetime, timedelta
from typing import Any

import httpx
from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse, RedirectResponse

_log = logging.getLogger(__name__)

auth_router = APIRouter(prefix="/auth", tags=["auth"])

_GOOGLE_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"
_GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
_GOOGLE_JWKS_URL = "https://www.googleapis.com/oauth2/v3/certs"
_REDIRECT_URI = "https://dashboard.complyform.dev/auth/callback"
_SCOPE = "openid email profile"
_STATE_EXPIRY_MINUTES = 10

_jwks_cache: dict[str, Any] | None = None
_jwks_cached_at: float = 0
_JWKS_CACHE_TTL = 86400  # 24 hours


def _get_oauth_client_id() -> str:
    """Load Google OAuth client ID from Secret Manager."""
    from google.cloud import secretmanager

    client = secretmanager.SecretManagerServiceClient()
    name = "projects/complyform-prod/secrets/google-oauth-client-id/versions/latest"
    response = client.access_secret_version(request={"name": name})
    return response.payload.data.decode()


def _get_oauth_client_secret() -> str:
    """Load Google OAuth client secret from Secret Manager."""
    from google.cloud import secretmanager

    client = secretmanager.SecretManagerServiceClient()
    name = "projects/complyform-prod/secrets/google-oauth-client-secret/versions/latest"
    response = client.access_secret_version(request={"name": name})
    return response.payload.data.decode()


async def _get_google_jwks() -> dict[str, Any]:
    """Fetch Google JWKS with 24-hour cache."""
    global _jwks_cache, _jwks_cached_at  # noqa: PLW0603
    now = time.monotonic()
    if _jwks_cache is not None and (now - _jwks_cached_at) < _JWKS_CACHE_TTL:
        return _jwks_cache

    async with httpx.AsyncClient() as client:
        resp = await client.get(_GOOGLE_JWKS_URL)
        resp.raise_for_status()
        _jwks_cache = resp.json()
        _jwks_cached_at = now
        return _jwks_cache  # type: ignore[return-value]


def _verify_id_token(
    id_token: str,
    jwks: dict[str, Any],
) -> dict[str, Any]:
    """Verify Google ID token and extract claims."""
    parts = id_token.split(".")
    if len(parts) != 3:
        msg = "Invalid JWT format"
        raise ValueError(msg)

    # Decode header
    header_b64 = parts[0] + "=" * (4 - len(parts[0]) % 4)
    header = json.loads(base64.b64decode(header_b64))
    kid = header.get("kid")

    # Find matching key
    matching_key = None
    for key in jwks.get("keys", []):
        if key.get("kid") == kid:
            matching_key = key
            break

    if matching_key is None:
        msg = "No matching key found in JWKS"
        raise ValueError(msg)

    # Decode payload (claims)
    payload_b64 = parts[1] + "=" * (4 - len(parts[1]) % 4)
    claims: dict[str, Any] = json.loads(base64.b64decode(payload_b64))

    # Verify basic claims
    valid_issuers = (
        "accounts.google.com",
        "https://accounts.google.com",
    )
    if claims.get("iss") not in valid_issuers:
        msg = "Invalid issuer"
        raise ValueError(msg)

    exp = claims.get("exp", 0)
    if datetime.fromtimestamp(exp, tz=UTC) < datetime.now(UTC):
        msg = "Token expired"
        raise ValueError(msg)

    return claims


async def _create_oauth_state() -> str:
    """Generate and store OAuth state token in Firestore."""
    from backend.shared.firestore import get_db

    state = secrets.token_hex(32)
    now = datetime.now(UTC)
    expires_at = now + timedelta(minutes=_STATE_EXPIRY_MINUTES)

    await (
        get_db()
        .collection("sessions")
        .document(state)
        .set(
            {
                "type": "oauth_state",
                "created_at": now.isoformat(),
                "expires_at": expires_at,
            }
        )
    )
    return state


def _build_google_redirect(state: str) -> RedirectResponse:
    """Build redirect URL to Google OAuth consent screen."""
    client_id = _get_oauth_client_id()
    params = {
        "client_id": client_id,
        "redirect_uri": _REDIRECT_URI,
        "response_type": "code",
        "scope": _SCOPE,
        "state": state,
        "prompt": "select_account",
    }
    query_string = "&".join(f"{k}={v}" for k, v in params.items())
    redirect_url = f"{_GOOGLE_AUTH_URL}?{query_string}"
    return RedirectResponse(url=redirect_url, status_code=302)


@auth_router.get("/login", response_model=None)
async def login(request: Request) -> RedirectResponse | JSONResponse:
    """Email-first login routing — SSO or Google OAuth.

    Query params:
    - provider=google → bypass email-first, go to Google OAuth
    - email=user@domain.com → check SSO config for domain
    """
    provider = request.query_params.get("provider", "")
    email = request.query_params.get("email", "")

    # Backward compat: explicit Google OAuth
    if provider == "google" or (not email and not provider):
        state = await _create_oauth_state()
        return _build_google_redirect(state)

    # Email-first: check SSO config
    domain = email.split("@")[-1] if "@" in email else ""
    if not domain:
        state = await _create_oauth_state()
        return _build_google_redirect(state)

    from backend.dashboard.sso import get_sso_config_by_domain

    sso_config = await get_sso_config_by_domain(domain)

    if sso_config is None:
        # No SSO → Google OAuth
        state = await _create_oauth_state()
        return _build_google_redirect(state)

    if sso_config.enforce_sso:
        # SSO enforced → redirect directly to IdP
        from backend.dashboard.sso import initiate_sso_login

        state = await _create_oauth_state()
        redirect_url = await initiate_sso_login(domain, state)
        return RedirectResponse(url=redirect_url, status_code=302)

    # SSO available but not enforced → for now, redirect to SSO
    # (in production, render choice page with htmx)
    from backend.dashboard.sso import initiate_sso_login

    state = await _create_oauth_state()
    redirect_url = await initiate_sso_login(domain, state)
    return RedirectResponse(url=redirect_url, status_code=302)


@auth_router.get("/callback", response_model=None)
async def callback(
    request: Request,
    code: str = "",
    state: str = "",
) -> RedirectResponse | JSONResponse:
    """Handle Google OAuth callback — exchange code, create session."""
    from backend.dashboard.session import (
        create_session,
        get_cookie_max_age,
        get_cookie_name,
        write_session,
    )
    from backend.shared.firestore import get_db

    # Verify state
    if not state:
        return RedirectResponse(url="/auth/login", status_code=302)

    state_doc = await get_db().collection("sessions").document(state).get()
    if not state_doc.exists:
        return RedirectResponse(url="/auth/login", status_code=302)

    state_data: dict[str, Any] = state_doc.to_dict()  # type: ignore[assignment]
    if state_data.get("type") != "oauth_state":
        return RedirectResponse(url="/auth/login", status_code=302)

    # Check expiry
    expires_at = state_data.get("expires_at")
    if expires_at is not None:
        now = datetime.now(UTC)
        if hasattr(expires_at, "timestamp"):
            exp_dt = (
                expires_at.replace(tzinfo=UTC)
                if expires_at.tzinfo is None
                else expires_at
            )
            if exp_dt < now:
                await get_db().collection("sessions").document(state).delete()
                return RedirectResponse(
                    url="/auth/login",
                    status_code=302,
                )
        elif isinstance(expires_at, str):
            try:
                exp_dt = datetime.fromisoformat(expires_at)
                if exp_dt < now:
                    await get_db().collection("sessions").document(state).delete()
                    return RedirectResponse(
                        url="/auth/login",
                        status_code=302,
                    )
            except ValueError, TypeError:
                pass

    # Delete used state
    await get_db().collection("sessions").document(state).delete()

    if not code:
        return RedirectResponse(url="/auth/login", status_code=302)

    # Exchange code for tokens
    try:
        client_id = _get_oauth_client_id()
        client_secret = _get_oauth_client_secret()

        async with httpx.AsyncClient() as http_client:
            token_resp = await http_client.post(
                _GOOGLE_TOKEN_URL,
                data={
                    "code": code,
                    "client_id": client_id,
                    "client_secret": client_secret,
                    "redirect_uri": _REDIRECT_URI,
                    "grant_type": "authorization_code",
                },
            )
            token_resp.raise_for_status()
            token_data = token_resp.json()
    except Exception:
        _log.exception("Token exchange failed")
        return JSONResponse(
            {
                "error": "Authentication failed. Please try again.",
                "retry_url": "/auth/login",
            },
            status_code=500,
        )

    # Verify ID token
    id_token = token_data.get("id_token", "")
    if not id_token:
        return JSONResponse(
            {
                "error": "No ID token received",
                "retry_url": "/auth/login",
            },
            status_code=500,
        )

    try:
        jwks = await _get_google_jwks()
        claims = _verify_id_token(id_token, jwks)
    except Exception:
        _log.exception("ID token verification failed")
        return JSONResponse(
            {
                "error": "Token verification failed",
                "retry_url": "/auth/login",
            },
            status_code=500,
        )

    email = claims.get("email", "")
    if not email:
        return JSONResponse(
            {
                "error": "No email in token claims",
                "retry_url": "/auth/login",
            },
            status_code=500,
        )

    # Find customer
    db = get_db()
    query = db.collection("customers").where("email", "==", email).limit(1)
    docs = [doc async for doc in query.stream()]

    if not docs:
        customer_id = ""
        tier = "community"
    else:
        customer_id = docs[0].id
        cust_data: dict[str, Any] = docs[0].to_dict()  # type: ignore[assignment]
        tier = cust_data.get("tier", "community")

    # Create session
    session_id, encrypted_value, session_doc, _csrf = create_session(
        email=email,
        customer_id=customer_id,
        tier=tier,
        request=request,
    )
    await write_session(session_id, session_doc)

    # Audit: dashboard.login (non-fatal)
    if customer_id:
        try:
            import hashlib

            from backend.dashboard.audit import write_audit_event
            from backend.shared.tier_config import TIER_CONFIG

            _host = request.client.host if request.client else "unknown"
            _ip_hash = hashlib.sha256(_host.encode()).hexdigest()
            _retention = TIER_CONFIG.get(tier, {}).get("retention_months")
            await write_audit_event(
                customer_id=customer_id,
                event="dashboard.login",
                actor_email=email,
                actor_role="owner",
                ip_hash=_ip_hash,
                metadata={"auth_method": "google_oauth"},
                retention_months=_retention,
            )
        except Exception:
            _log.exception("Failed to write login audit event")

    # Set cookie and redirect
    response = RedirectResponse(url="/dashboard/", status_code=302)
    response.set_cookie(
        key=get_cookie_name(),
        value=encrypted_value,
        httponly=True,
        secure=True,
        samesite="lax",
        path="/",
        max_age=get_cookie_max_age(),
    )
    return response


@auth_router.get("/sso/callback", response_model=None)
async def sso_callback(
    request: Request,
    code: str = "",
    state: str = "",
) -> RedirectResponse | JSONResponse:
    """Handle SSO callback — exchange code, create session."""
    from backend.dashboard.session import (
        create_session,
        get_cookie_max_age,
        get_cookie_name,
        write_session,
    )
    from backend.shared.firestore import get_db

    # Verify state
    if not state:
        return RedirectResponse(url="/auth/login", status_code=302)

    state_doc = await get_db().collection("sessions").document(state).get()
    if not state_doc.exists:
        return RedirectResponse(url="/auth/login", status_code=302)

    state_data: dict[str, Any] = state_doc.to_dict()  # type: ignore[assignment]
    if state_data.get("type") != "oauth_state":
        return RedirectResponse(url="/auth/login", status_code=302)

    # Delete used state
    await get_db().collection("sessions").document(state).delete()

    if not code:
        return RedirectResponse(url="/auth/login", status_code=302)

    # Exchange code via WorkOS
    try:
        from backend.dashboard.sso import handle_sso_callback

        sso_user = await handle_sso_callback(code)
    except Exception:
        _log.exception("SSO callback failed")
        return JSONResponse(
            {
                "error": "SSO authentication failed. Please try again.",
                "retry_url": "/auth/login",
            },
            status_code=500,
        )

    email = sso_user.email

    # Find customer
    db = get_db()
    query = db.collection("customers").where("email", "==", email).limit(1)
    docs = [doc async for doc in query.stream()]

    if not docs:
        customer_id = ""
        tier = "community"
    else:
        customer_id = docs[0].id
        cust_data: dict[str, Any] = docs[0].to_dict()  # type: ignore[assignment]
        tier = cust_data.get("tier", "community")

    # Create session with SSO auth method
    session_id, encrypted_value, session_doc, _csrf = create_session(
        email=email,
        customer_id=customer_id,
        tier=tier,
        request=request,
    )
    # Add SSO-specific fields
    session_doc["auth_method"] = "saml_sso"
    session_doc["idp_connection_id"] = sso_user.connection_id
    await write_session(session_id, session_doc)

    # Audit: dashboard.login (non-fatal)
    if customer_id:
        try:
            import hashlib

            from backend.dashboard.audit import write_audit_event
            from backend.shared.tier_config import TIER_CONFIG

            _host = request.client.host if request.client else "unknown"
            _ip_hash = hashlib.sha256(_host.encode()).hexdigest()
            _retention = TIER_CONFIG.get(tier, {}).get("retention_months")
            await write_audit_event(
                customer_id=customer_id,
                event="dashboard.login",
                actor_email=email,
                actor_role="owner",
                ip_hash=_ip_hash,
                metadata={"auth_method": "saml_sso"},
                retention_months=_retention,
            )
        except Exception:
            _log.exception("Failed to write SSO login audit event")

    # Set cookie and redirect
    response = RedirectResponse(url="/dashboard/", status_code=302)
    response.set_cookie(
        key=get_cookie_name(),
        value=encrypted_value,
        httponly=True,
        secure=True,
        samesite="lax",
        path="/",
        max_age=get_cookie_max_age(),
    )
    return response


@auth_router.post("/logout", response_model=None)
async def logout(
    request: Request,
) -> RedirectResponse | JSONResponse:
    """Logout — validate CSRF, delete session, clear cookie."""
    from backend.dashboard.session import (
        delete_session,
        get_cookie_name,
        get_session,
        validate_csrf,
    )

    cookie_name = get_cookie_name()
    encrypted_cookie = request.cookies.get(cookie_name)
    if not encrypted_cookie:
        return RedirectResponse(
            url="https://complyform.dev",
            status_code=302,
        )

    session = await get_session(encrypted_cookie)
    if session is None:
        response = RedirectResponse(
            url="https://complyform.dev",
            status_code=302,
        )
        response.delete_cookie(cookie_name, path="/")
        return response

    # Validate CSRF
    csrf_from_header = request.headers.get("X-CSRF-Token", "")
    csrf_from_form = ""
    if not csrf_from_header:
        try:
            form = await request.form()
            csrf_from_form = str(form.get("_csrf", ""))
        except Exception:
            pass

    request_csrf = csrf_from_header or csrf_from_form
    session_csrf = session.get("csrf_token", "")

    if not request_csrf or not validate_csrf(request_csrf, session_csrf):
        return JSONResponse(
            {"error": "CSRF validation failed"},
            status_code=403,
        )

    # Delete session
    session_id = session.get("session_id", "")
    if session_id:
        await delete_session(session_id)

    # Clear cookie and redirect
    response = RedirectResponse(
        url="https://complyform.dev",
        status_code=302,
    )
    response.set_cookie(
        key=cookie_name,
        value="",
        httponly=True,
        secure=True,
        samesite="lax",
        path="/",
        max_age=0,
    )
    return response
