//! Delegated login: "Sign in with Makenot.work" (OAuth client side).
//!
//! Used on the testnot.work staging mirror. Instead of a local password form,
//! the login page redirects to an upstream MNW provider (production), where the
//! user authenticates — so a password is only ever entered on the real site.
//! On callback we exchange the code for the provider's response, take the
//! verified `user_id`, look that user up in our own (mirrored) DB, and start a
//! local session. The provider's OAuth flow is the SyncKit one
//! (`src/routes/oauth.rs`); we discard its sync token and use only `user_id`.
//!
//! Active only when `[sso]` is configured (the three `SSO_*` vars). Routes are
//! allowlisted in the access gate so an unauthenticated visitor can reach them.

use axum::{
    extract::{Query, State},
    http::HeaderMap,
    response::{IntoResponse, Redirect, Response},
    routing::get,
    Router,
};
use base64::Engine;
use rand::RngCore;
use serde::Deserialize;
use sha2::{Digest, Sha256};
use tower_sessions::Session;

use crate::{
    auth::{login_user, track_session, SessionUser},
    db::{self, UserId},
    error::{AppError, Result},
    AppState,
};

const SSO_STATE_KEY: &str = "sso_state";
const SSO_VERIFIER_KEY: &str = "sso_pkce_verifier";

fn b64url(data: &[u8]) -> String {
    base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(data)
}

/// GET /sso/login — begin the delegated-login flow.
///
/// Generates PKCE + state, stashes them in the session, and redirects to the
/// provider's authorize endpoint. No-op (404-ish redirect home) when SSO is off.
#[tracing::instrument(skip_all, name = "sso::login")]
async fn sso_login(State(state): State<AppState>, session: Session) -> Result<Response> {
    let Some(sso) = state.config.sso.as_ref() else {
        // SSO not configured — nothing to delegate to.
        return Ok(Redirect::to("/login").into_response());
    };

    // PKCE verifier (43-char base64url of 32 random bytes) + S256 challenge.
    let mut vbytes = [0u8; 32];
    rand::rng().fill_bytes(&mut vbytes);
    let verifier = b64url(&vbytes);
    let challenge = b64url(Sha256::digest(verifier.as_bytes()).as_ref());

    // CSRF-style state to bind the callback to this session.
    let mut sbytes = [0u8; 16];
    rand::rng().fill_bytes(&mut sbytes);
    let state_param = b64url(&sbytes);

    session.insert(SSO_STATE_KEY, &state_param).await.map_err(|e| AppError::Internal(e.into()))?;
    session.insert(SSO_VERIFIER_KEY, &verifier).await.map_err(|e| AppError::Internal(e.into()))?;

    let redirect_uri = format!("{}/sso/callback", state.config.host_url);
    let authorize = format!(
        "{}/oauth/authorize?response_type=code&client_id={}&redirect_uri={}&state={}&code_challenge={}&code_challenge_method=S256",
        sso.provider_url,
        urlencoding::encode(&sso.client_id),
        urlencoding::encode(&redirect_uri),
        urlencoding::encode(&state_param),
        urlencoding::encode(&challenge),
    );
    Ok(Redirect::to(&authorize).into_response())
}

#[derive(Deserialize)]
struct CallbackQuery {
    code: Option<String>,
    state: Option<String>,
    error: Option<String>,
}

/// Minimal view of the provider's token response — we only need the user id.
#[derive(Deserialize)]
struct TokenResponse {
    user_id: UserId,
}

/// GET /sso/callback — provider redirected back with `code` + `state`.
#[tracing::instrument(skip_all, name = "sso::callback")]
async fn sso_callback(
    State(state): State<AppState>,
    session: Session,
    headers: HeaderMap,
    Query(q): Query<CallbackQuery>,
) -> Result<Response> {
    let Some(sso) = state.config.sso.as_ref() else {
        return Ok(Redirect::to("/login").into_response());
    };

    let fail = |msg: &str| Ok(Redirect::to(&format!("/login?sso_error={}", urlencoding::encode(msg))).into_response());

    if let Some(err) = q.error.as_deref() {
        tracing::warn!(error = %err, "sso provider returned error");
        return fail("Sign-in was cancelled or denied.");
    }
    let (Some(code), Some(returned_state)) = (q.code.as_deref(), q.state.as_deref()) else {
        return fail("Sign-in response was incomplete. Please try again.");
    };

    // Validate state against the session, and consume the one-shot PKCE values.
    let expected_state: Option<String> = session.get(SSO_STATE_KEY).await.ok().flatten();
    let verifier: Option<String> = session.get(SSO_VERIFIER_KEY).await.ok().flatten();
    let _ = session.remove::<String>(SSO_STATE_KEY).await;
    let _ = session.remove::<String>(SSO_VERIFIER_KEY).await;

    let (Some(expected_state), Some(verifier)) = (expected_state, verifier) else {
        return fail("Your sign-in session expired. Please try again.");
    };
    if !crate::helpers::constant_time_compare(&expected_state, returned_state) {
        tracing::warn!("sso state mismatch");
        return fail("Sign-in could not be verified. Please try again.");
    }

    // Exchange the code at the provider's token endpoint.
    let redirect_uri = format!("{}/sso/callback", state.config.host_url);
    let resp = reqwest::Client::new()
        .post(format!("{}/oauth/token", sso.provider_url))
        .timeout(std::time::Duration::from_secs(10))
        .form(&[
            ("grant_type", "authorization_code"),
            ("code", code),
            ("redirect_uri", &redirect_uri),
            ("code_verifier", &verifier),
            ("client_id", &sso.client_id),
            ("key", &sso.key),
        ])
        .send()
        .await;

    let resp = match resp {
        Ok(r) if r.status().is_success() => r,
        Ok(r) => {
            tracing::warn!(status = %r.status(), "sso token exchange rejected");
            return fail("Sign-in failed at the provider. Please try again.");
        }
        Err(e) => {
            tracing::warn!(error = ?e, "sso token exchange request failed");
            return fail("Could not reach the sign-in provider. Please try again.");
        }
    };

    let token: TokenResponse = match resp.json().await {
        Ok(t) => t,
        Err(e) => {
            tracing::warn!(error = ?e, "sso token response parse failed");
            return fail("Sign-in failed at the provider. Please try again.");
        }
    };

    // Map the verified provider user id onto our mirrored account.
    let db_user = match db::users::get_user_by_id(&state.db, token.user_id).await? {
        Some(u) => u,
        None => return fail("Your account isn't in the preview yet — it syncs from production daily."),
    };
    if db_user.is_suspended() || db_user.is_deactivated() {
        return fail("This account is not active.");
    }

    let user_id = db_user.id;
    let session_user = SessionUser::from_db_user(db_user, &state.db, state.config.admin_user_id).await;
    login_user(&session, session_user).await?;
    track_session(&session, &state.db, user_id, &headers).await?;

    Ok(Redirect::to("/").into_response())
}

pub fn sso_routes() -> Router<AppState> {
    Router::new()
        .route("/sso/login", get(sso_login))
        .route("/sso/callback", get(sso_callback))
}