"""SSO integration — WorkOS SAML/OIDC for dashboard authentication.

Manages SSO configurations per customer, initiates IdP login flows,
and handles callbacks via WorkOS Python SDK.
"""

from __future__ import annotations

import logging
from datetime import UTC, datetime
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict

_log = logging.getLogger(__name__)


class SSOConfig(BaseModel, frozen=True):
    """Firestore sso_configs/{customer_id} document."""

    model_config = ConfigDict(strict=True)
    customer_id: str
    enabled: bool
    broker_connection_id: str
    email_domain: str
    enforce_sso: bool
    provider_type: Literal["saml", "oidc"]
    provider_name: Literal["okta", "entra_id", "google_workspace", "ping", "custom"]
    created_at: str
    updated_at: str
    created_by: str
    group_mapping: dict[str, str] = {}


class SSOConfigInput(BaseModel, frozen=True):
    """Input model for creating/updating SSO configuration."""

    model_config = ConfigDict(strict=True)
    enabled: bool
    email_domain: str
    enforce_sso: bool
    provider_type: Literal["saml", "oidc"]
    provider_name: Literal["okta", "entra_id", "google_workspace", "ping", "custom"]
    group_mapping: dict[str, str] = {}


class SSOUser(BaseModel, frozen=True):
    """User profile returned from SSO callback."""

    model_config = ConfigDict(strict=True)
    email: str
    first_name: str | None = None
    last_name: str | None = None
    idp_id: str
    connection_id: str


def _get_workos_api_key() -> str:
    """Load WorkOS API key from Secret Manager."""
    from google.cloud import secretmanager

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


def _get_workos_client_id() -> str:
    """Load WorkOS client ID from Secret Manager."""
    from google.cloud import secretmanager

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


async def get_sso_config(customer_id: str) -> SSOConfig | None:
    """Load SSO config from Firestore sso_configs/{customer_id}."""
    from backend.shared.firestore import get_db

    doc = await get_db().collection("sso_configs").document(customer_id).get()
    if not doc.exists:
        return None

    data: dict[str, Any] = doc.to_dict()  # type: ignore[assignment]
    return SSOConfig(**data)


async def get_sso_config_by_domain(
    email_domain: str,
) -> SSOConfig | None:
    """Look up enabled SSO config by email domain."""
    from backend.shared.firestore import get_db

    query = (
        get_db()
        .collection("sso_configs")
        .where("email_domain", "==", email_domain)
        .where("enabled", "==", True)
        .limit(1)
    )
    docs = [doc async for doc in query.stream()]
    if not docs:
        return None

    data: dict[str, Any] = docs[0].to_dict()  # type: ignore[assignment]
    return SSOConfig(**data)


async def initiate_sso_login(email_domain: str, state: str) -> str:
    """Initiate SSO login flow via WorkOS.

    Returns redirect URL to the IdP.
    Raises SSONotConfiguredError if no config found.
    """
    from src.complyform.core.errors import SSONotConfiguredError

    config = await get_sso_config_by_domain(email_domain)
    if config is None:
        raise SSONotConfiguredError(email_domain)

    import workos

    workos_client = workos.WorkOSClient(
        api_key=_get_workos_api_key(),
        client_id=_get_workos_client_id(),
    )

    auth_url = workos_client.sso.get_authorization_url(
        connection_id=config.broker_connection_id,
        redirect_uri="https://dashboard.complyform.dev/auth/sso/callback",
        state=state,
    )

    return str(auth_url)


async def handle_sso_callback(code: str) -> SSOUser:
    """Exchange authorization code for user profile via WorkOS.

    Returns SSOUser with identity details.
    Raises SSOAuthenticationFailedError on failure.
    """
    from src.complyform.core.errors import SSOAuthenticationFailedError

    try:
        import workos

        workos_client = workos.WorkOSClient(
            api_key=_get_workos_api_key(),
            client_id=_get_workos_client_id(),
        )

        profile_and_token = workos_client.sso.get_profile_and_token(code)
        profile = profile_and_token.profile

        return SSOUser(
            email=profile.email,
            first_name=profile.first_name,
            last_name=profile.last_name,
            idp_id=profile.idp_id,
            connection_id=profile.connection_id,
        )
    except Exception as exc:
        _log.exception("SSO callback failed")
        raise SSOAuthenticationFailedError(str(exc)) from exc


async def save_sso_config(
    customer_id: str,
    config: SSOConfigInput,
    actor_email: str,
    connection_id: str,
) -> None:
    """Write SSO config to Firestore sso_configs/{customer_id}."""
    from backend.shared.firestore import get_db

    now = datetime.now(UTC).isoformat()

    existing = await get_sso_config(customer_id)
    created_at = existing.created_at if existing else now

    doc_data = {
        "customer_id": customer_id,
        "enabled": config.enabled,
        "broker_connection_id": connection_id,
        "email_domain": config.email_domain,
        "enforce_sso": config.enforce_sso,
        "provider_type": config.provider_type,
        "provider_name": config.provider_name,
        "created_at": created_at,
        "updated_at": now,
        "created_by": actor_email,
        "group_mapping": config.group_mapping,
    }

    await get_db().collection("sso_configs").document(customer_id).set(doc_data)


async def delete_sso_config(customer_id: str) -> None:
    """Delete SSO config from Firestore."""
    from backend.shared.firestore import get_db

    await get_db().collection("sso_configs").document(customer_id).delete()
