"""Single sign-on through a tenant's own identity provider, over OpenID
Connect — Microsoft Entra ID, Okta, Google Workspace, Keycloak, Auth0 and
anything else that speaks it.

The authorization-code flow with PKCE:
  start()   remembers the attempt (sso_states: a nonce and a PKCE verifier,
            used once) and sends the browser to the provider;
  finish()  takes the provider's answer: the state must be a live attempt
            for this tenant, the code is exchanged (with the verifier and
            the client secret) for an ID token, and the token must be signed
            by one of the provider's published keys, issued by the
            configured issuer, for our client ID, unexpired, and carry the
            attempt's nonce. Then the person is found — by the provider's
            subject id once linked, else by email (which links it) — or,
            when the tenant allows it, created.

The provider's addresses come from its discovery document; the issuer is
checked like a data query's URL (no link-local or cloud-metadata
addresses; the server's own network only with DATA_QUERY_ALLOW_PRIVATE_HOSTS).
"""

import base64
import hashlib
import secrets
import time
import uuid
from datetime import datetime, timedelta, timezone
from urllib.parse import urlencode, urlsplit

import httpx
from jose import JWTError, jwt
from sqlalchemy import select, text
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.tls import TRUSTSTORE_SSL_CONTEXT
from app.crud import licensing
from app.crud.sign_in_policy import SecuritySettings, client_secret
from app.models.identity import Role, Tenant, User, UserRole

STATE_TTL = timedelta(minutes=10)
_DISCOVERY_TTL = 3600
_discovery: dict[str, tuple[float, dict]] = {}
_jwks: dict[str, tuple[float, dict]] = {}


class OidcError(Exception):
    """Something the person should be told, on the sign-in page."""


def _http() -> httpx.AsyncClient:
    return httpx.AsyncClient(timeout=10, verify=TRUSTSTORE_SSL_CONTEXT, follow_redirects=False)


def _check_address(url: str) -> None:
    from app.crud.data_queries import DataQueryError, _check_host

    try:
        _check_host(url)
    except DataQueryError as exc:
        raise OidcError(f"The identity provider's address can't be used: {exc}") from exc


async def _get_json(url: str, cache: dict[str, tuple[float, dict]]) -> dict:
    hit = cache.get(url)
    if hit and hit[0] > time.monotonic():
        return hit[1]
    _check_address(url)
    try:
        async with _http() as client:
            resp = await client.get(url)
    except httpx.HTTPError as exc:
        raise OidcError(f"Couldn't reach the identity provider ({exc}).") from exc
    if resp.status_code != 200:
        raise OidcError(f"The identity provider answered HTTP {resp.status_code} for {url}.")
    doc = resp.json()
    cache[url] = (time.monotonic() + _DISCOVERY_TTL, doc)
    return doc


async def discover(issuer: str) -> dict:
    doc = await _get_json(f"{issuer}/.well-known/openid-configuration", _discovery)
    for key in ("authorization_endpoint", "token_endpoint", "jwks_uri"):
        if not doc.get(key):
            raise OidcError(f"The identity provider's discovery document has no {key}.")
    if doc.get("issuer", "").rstrip("/") != issuer:
        raise OidcError("The identity provider's discovery document names a different issuer.")
    return doc


def _pkce() -> tuple[str, str]:
    verifier = secrets.token_urlsafe(48)
    challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).decode().rstrip("=")
    return verifier, challenge


async def start(db: AsyncSession, tenant: Tenant, policy: SecuritySettings, redirect_uri: str, return_to: str) -> str:
    """The provider URL to send the browser to."""
    sso = policy.sso
    if not sso.enabled:
        raise OidcError("Single sign-on isn't set up for this workspace.")
    doc = await discover(sso.issuer)
    verifier, challenge = _pkce()
    nonce = secrets.token_urlsafe(24)
    state_id = (
        await db.execute(
            text("INSERT INTO sso_states (tenant_id, nonce, code_verifier, return_to) VALUES (:t, :n, :v, :r) RETURNING id"),
            {"t": tenant.id, "n": nonce, "v": verifier, "r": return_to},
        )
    ).scalar_one()
    # Old attempts nobody finished.
    await db.execute(text("DELETE FROM sso_states WHERE created_at < now() - make_interval(secs => :s)"), {"s": STATE_TTL.total_seconds()})
    query = {
        "response_type": "code",
        "client_id": sso.client_id,
        "redirect_uri": redirect_uri,
        "scope": "openid email profile",
        "state": str(state_id),
        "nonce": nonce,
        "code_challenge": challenge,
        "code_challenge_method": "S256",
    }
    sep = "&" if urlsplit(doc["authorization_endpoint"]).query else "?"
    return doc["authorization_endpoint"] + sep + urlencode(query)


async def take_state(db: AsyncSession, state: str) -> dict:
    """The attempt `state` names — used up by this, whatever happens next."""
    try:
        state_id = uuid.UUID(state)
    except (ValueError, TypeError):
        raise OidcError("This sign-in link isn't valid — start again.") from None
    row = (
        await db.execute(
            text("DELETE FROM sso_states WHERE id = :id RETURNING tenant_id, nonce, code_verifier, return_to, created_at"), {"id": state_id}
        )
    ).first()
    if row is None:
        raise OidcError("This sign-in attempt was already used or has expired — start again.")
    if row.created_at < datetime.now(timezone.utc) - STATE_TTL:
        raise OidcError("This sign-in attempt took too long — start again.")
    return dict(row._mapping)


async def _verify_id_token(id_token: str, access_token: str | None, policy: SecuritySettings, doc: dict, nonce: str) -> dict:
    sso = policy.sso
    try:
        header = jwt.get_unverified_header(id_token)
    except JWTError as exc:
        raise OidcError("The identity provider sent an unreadable ID token.") from exc
    if header.get("alg") in (None, "none") or header["alg"].startswith("HS"):
        raise OidcError("The identity provider's ID token isn't signed with a public key.")
    jwks = await _get_json(doc["jwks_uri"], _jwks)
    keys = [k for k in jwks.get("keys", []) if not header.get("kid") or k.get("kid") == header.get("kid")]
    if not keys:
        # A rotated key: fetch the set again once.
        _jwks.pop(doc["jwks_uri"], None)
        jwks = await _get_json(doc["jwks_uri"], _jwks)
        keys = [k for k in jwks.get("keys", []) if not header.get("kid") or k.get("kid") == header.get("kid")]
    if not keys:
        raise OidcError("The ID token was signed with a key the identity provider doesn't publish.")
    try:
        claims = jwt.decode(
            id_token, {"keys": keys}, algorithms=[header["alg"]], audience=sso.client_id, issuer=sso.issuer,
            access_token=access_token, options={"verify_at_hash": bool(access_token)},
        )
    except JWTError as exc:
        raise OidcError(f"The identity provider's ID token didn't check out: {exc}.") from exc
    if claims.get("nonce") != nonce:
        raise OidcError("The identity provider's answer doesn't belong to this sign-in attempt.")
    return claims


async def exchange(policy: SecuritySettings, code: str, verifier: str, redirect_uri: str, nonce: str) -> dict:
    """The verified ID token claims for an authorization code."""
    sso = policy.sso
    doc = await discover(sso.issuer)
    _check_address(doc["token_endpoint"])
    form = {
        "grant_type": "authorization_code", "code": code, "redirect_uri": redirect_uri, "code_verifier": verifier,
        "client_id": sso.client_id, "client_secret": client_secret(policy),
    }
    try:
        async with _http() as client:
            resp = await client.post(doc["token_endpoint"], data=form, headers={"Accept": "application/json"})
    except httpx.HTTPError as exc:
        raise OidcError(f"Couldn't reach the identity provider ({exc}).") from exc
    body = resp.json() if resp.headers.get("content-type", "").startswith("application/json") else {}
    if resp.status_code != 200 or not body.get("id_token"):
        reason = body.get("error_description") or body.get("error") or f"HTTP {resp.status_code}"
        raise OidcError(f"The identity provider didn't accept the sign-in: {reason}.")
    return await _verify_id_token(body["id_token"], body.get("access_token"), policy, doc, nonce)


def _email_of(claims: dict) -> str:
    email = claims.get("email") or (claims.get("preferred_username") if "@" in str(claims.get("preferred_username", "")) else None)
    if not email:
        raise OidcError("The identity provider didn't share an email address — allow the 'email' scope for this app.")
    if claims.get("email_verified") is False:
        raise OidcError("The identity provider hasn't verified this email address.")
    return str(email).strip().lower()


async def resolve_user(db: AsyncSession, tenant: Tenant, policy: SecuritySettings, claims: dict) -> User:
    """The account this provider identity signs in as."""
    sso = policy.sso
    subject = f"oidc:{claims['sub']}"
    email = _email_of(claims)
    domain = email.rsplit("@", 1)[-1]
    if sso.allowed_domains and domain not in sso.allowed_domains:
        raise OidcError(f"Accounts from {domain} can't sign in to this workspace.")

    user = (await db.execute(select(User).where(User.tenant_id == tenant.id, User.external_auth_id == subject))).scalar_one_or_none()
    if user is None:
        user = (await db.execute(select(User).where(User.tenant_id == tenant.id, User.email == email))).scalar_one_or_none()
        if user is not None:
            if user.external_auth_id and user.external_auth_id != subject:
                raise OidcError("This account is already linked to a different identity at the provider.")
            user.external_auth_id = subject
    if user is None:
        if not sso.auto_create_users:
            raise OidcError(f"There's no account for {email} in this workspace — ask an administrator to add you.")
        role = await db.get(Role, uuid.UUID(sso.default_role_id)) if sso.default_role_id else None
        if role is None or role.tenant_id != tenant.id or role.scope != "internal":
            raise OidcError("New accounts can't be created — the role for them no longer exists. Ask an administrator.")
        try:
            await licensing.require_seat(db, tenant.id, portal=False)
        except licensing.LicenseError as exc:
            raise OidcError(exc.message) from None
        name = claims.get("name") or " ".join(p for p in (claims.get("given_name"), claims.get("family_name")) if p) or email
        user = User(tenant_id=tenant.id, name=str(name)[:200], email=email, external_auth_id=subject, status="active")
        db.add(user)
        await db.flush()
        db.add(UserRole(user_id=user.id, role_id=role.id))
    if user.status != "active":
        raise OidcError("This account is inactive.")
    return user