//! OIDC authorization-code + PKCE flow against Kanidm.
//!
//! Kanidm serves per-client OIDC discovery documents at
//! `<KANIDM_URL>/oauth2/openid/<client_id>/.well-known/openid-configuration`,
//! and enforces PKCE, so this module always sends a S256 challenge.

use axum::{
    extract::{Query, State},
    http::StatusCode,
    response::Redirect,
};
use openidconnect::{
    core::{CoreAuthenticationFlow, CoreClient, CoreProviderMetadata},
    AuthorizationCode, ClientId, ClientSecret, CsrfToken, EndpointMaybeSet, EndpointNotSet,
    EndpointSet, IssuerUrl, Nonce, PkceCodeChallenge, PkceCodeVerifier, RedirectUrl, Scope,
    TokenResponse,
};
use serde::Deserialize;
use tower_sessions::Session;

use crate::auth::{User, SESSION_USER_KEY};

use super::AppState;

type OidcClient = CoreClient<
    EndpointSet,      // auth endpoint
    EndpointNotSet,   // device auth
    EndpointNotSet,   // introspection
    EndpointNotSet,   // revocation
    EndpointMaybeSet, // token endpoint (from discovery)
    EndpointMaybeSet, // userinfo endpoint (from discovery)
>;

pub struct Oidc {
    client: OidcClient,
    http: openidconnect::reqwest::Client,
}

const PKCE_KEY: &str = "oidc_pkce_verifier";
const CSRF_KEY: &str = "oidc_csrf_state";
const NONCE_KEY: &str = "oidc_nonce";

impl Oidc {
    /// Discovers the provider and builds the client from environment:
    /// `KANIDM_URL`, `OAUTH2_CLIENT_ID`, `OAUTH2_CLIENT_SECRET`, `PUBLIC_URL`.
    pub async fn from_env() -> anyhow::Result<Self> {
        let kanidm_url = require_env("KANIDM_URL")?;
        let client_id = require_env("OAUTH2_CLIENT_ID")?;
        let client_secret = require_env("OAUTH2_CLIENT_SECRET")?;
        let public_url = require_env("PUBLIC_URL")?;

        let issuer = IssuerUrl::new(format!(
            "{}/oauth2/openid/{}",
            kanidm_url.trim_end_matches('/'),
            client_id
        ))?;
        let redirect = RedirectUrl::new(format!(
            "{}/auth/callback",
            public_url.trim_end_matches('/')
        ))?;

        // Never follow redirects when talking to the IdP (SSRF hygiene).
        let http = openidconnect::reqwest::ClientBuilder::new()
            .redirect(openidconnect::reqwest::redirect::Policy::none())
            .build()?;

        tracing::info!(issuer = %issuer.as_str(), "discovering OIDC provider");
        let metadata = CoreProviderMetadata::discover_async(issuer, &http).await?;
        let client = CoreClient::from_provider_metadata(
            metadata,
            ClientId::new(client_id),
            Some(ClientSecret::new(client_secret)),
        )
        .set_redirect_uri(redirect);

        Ok(Self { client, http })
    }
}

fn require_env(name: &str) -> anyhow::Result<String> {
    std::env::var(name).map_err(|_| anyhow::anyhow!("missing required env var {name}"))
}

type HandlerError = (StatusCode, String);

fn internal(err: impl std::fmt::Display) -> HandlerError {
    tracing::error!("oidc error: {err}");
    (
        StatusCode::INTERNAL_SERVER_ERROR,
        "authentication failed; see server logs".to_string(),
    )
}

/// GET /auth/login — stash PKCE/state/nonce in the session and bounce to Kanidm.
pub async fn login(State(state): State<AppState>, session: Session) -> Result<Redirect, HandlerError> {
    let (pkce_challenge, pkce_verifier) = PkceCodeChallenge::new_random_sha256();

    let (auth_url, csrf_state, nonce) = state
        .oidc
        .client
        .authorize_url(
            CoreAuthenticationFlow::AuthorizationCode,
            CsrfToken::new_random,
            Nonce::new_random,
        )
        .add_scope(Scope::new("openid".to_string()))
        .add_scope(Scope::new("profile".to_string()))
        .add_scope(Scope::new("email".to_string()))
        .set_pkce_challenge(pkce_challenge)
        .url();

    session
        .insert(PKCE_KEY, pkce_verifier.secret())
        .await
        .map_err(internal)?;
    session
        .insert(CSRF_KEY, csrf_state.secret())
        .await
        .map_err(internal)?;
    session
        .insert(NONCE_KEY, nonce.secret())
        .await
        .map_err(internal)?;

    Ok(Redirect::to(auth_url.as_str()))
}

#[derive(Deserialize)]
pub struct CallbackParams {
    code: String,
    state: String,
}

/// GET /auth/callback — verify state, exchange the code, verify the ID token,
/// and store the user in the session.
pub async fn callback(
    State(state): State<AppState>,
    session: Session,
    Query(params): Query<CallbackParams>,
) -> Result<Redirect, HandlerError> {
    let stored_csrf: Option<String> = session.remove(CSRF_KEY).await.map_err(internal)?;
    let pkce_verifier: Option<String> = session.remove(PKCE_KEY).await.map_err(internal)?;
    let nonce: Option<String> = session.remove(NONCE_KEY).await.map_err(internal)?;

    let (Some(stored_csrf), Some(pkce_verifier), Some(nonce)) =
        (stored_csrf, pkce_verifier, nonce)
    else {
        return Err((
            StatusCode::BAD_REQUEST,
            "no login in progress; start again at /auth/login".to_string(),
        ));
    };
    if params.state != stored_csrf {
        return Err((
            StatusCode::BAD_REQUEST,
            "state mismatch; start again at /auth/login".to_string(),
        ));
    }

    let oidc = &state.oidc;
    let token_response = oidc
        .client
        .exchange_code(AuthorizationCode::new(params.code))
        .map_err(internal)?
        .set_pkce_verifier(PkceCodeVerifier::new(pkce_verifier))
        .request_async(&oidc.http)
        .await
        .map_err(internal)?;

    let id_token = token_response
        .id_token()
        .ok_or_else(|| internal("provider returned no ID token"))?;
    let claims = id_token
        .claims(&oidc.client.id_token_verifier(), &Nonce::new(nonce))
        .map_err(internal)?;

    let username = claims
        .preferred_username()
        .map(|u| u.as_str().to_string())
        .or_else(|| claims.email().map(|e| e.as_str().to_string()))
        .unwrap_or_else(|| claims.subject().as_str().to_string());
    let display_name = claims
        .name()
        .and_then(|n| n.get(None))
        .map(|n| n.as_str().to_string())
        .unwrap_or_else(|| username.clone());

    // `groups` is a custom claim (Kanidm `oauth2 update-claim-map`), not
    // something the Core* typed claims struct above knows about. The
    // signature is already verified by `id_token.claims(...)` above, so
    // re-reading the same payload's raw JSON for one more field is safe -
    // just a plain field extraction, not a second verification step.
    // IdToken's Serialize impl (not Display - it has none) produces the
    // raw compact JWT string "header.payload.signature".
    let groups = extract_groups_claim(&id_token);

    let user = User {
        sub: claims.subject().as_str().to_string(),
        username,
        display_name,
        groups,
    };

    // Rotate the session id on privilege change, then store the user.
    session.cycle_id().await.map_err(internal)?;
    session
        .insert(SESSION_USER_KEY, &user)
        .await
        .map_err(internal)?;

    tracing::info!(user = %user.username, "signed in");
    Ok(Redirect::to("/"))
}

/// GET /auth/logout — drop the session.
pub async fn logout(session: Session) -> Result<Redirect, HandlerError> {
    session.flush().await.map_err(internal)?;
    Ok(Redirect::to("/"))
}

/// Pulls the `groups` custom claim (Kanidm `oauth2 update-claim-map`) out
/// of an ID token's raw JWT payload. `IdToken`'s `Serialize` impl (it has
/// no `Display`) produces the compact "header.payload.signature" string,
/// which is where this reads from - the signature itself is never
/// re-checked here, that already happened via `id_token.claims(...)`
/// before this is called. Defensive by design: any parse failure (no
/// claim, wrong shape) just yields no groups rather than failing login.
fn extract_groups_claim<T: serde::Serialize>(id_token: &T) -> Vec<String> {
    use base64::Engine;

    let Ok(serde_json::Value::String(compact)) = serde_json::to_value(id_token) else {
        return Vec::new();
    };
    let Some(payload_b64) = compact.split('.').nth(1) else {
        return Vec::new();
    };
    let Ok(payload_bytes) = base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(payload_b64)
    else {
        return Vec::new();
    };
    let Ok(payload) = serde_json::from_slice::<serde_json::Value>(&payload_bytes) else {
        return Vec::new();
    };

    payload
        .get("groups")
        .and_then(|g| g.as_array())
        .map(|arr| arr.iter().filter_map(|v| v.as_str().map(String::from)).collect())
        .unwrap_or_default()
}