//! OIDC and OAuth 2.0 protocol endpoints.

use crate::errors::HttpError;
use crate::handlers::AppStateExt;
use axum::extract::Query;
use axum::http::{header, HeaderMap};
use axum::response::{IntoResponse, Redirect, Response};
use axum::Form;
use axum::Json;
use base64ct::{Base64, Encoding};
use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
use serde::{Deserialize, Serialize};
use std::str::FromStr;
use sui_id_core::authorize::{
    self, AuthorizeParams, CodeExchangeRequest, IssuanceContext, RefreshExchangeRequest,
};
use sui_id_core::discovery::Discovery;
use sui_id_core::errors::{CoreError, ProtocolError};
use sui_id_core::jwks;
use sui_id_shared::ids::ClientId;
use sui_id_store::repos::users;

// ---------- discovery & JWKS ----------

pub async fn discovery(state_ext: AppStateExt) -> impl IntoResponse {
    let axum::extract::State(app) = state_ext;
    Json(Discovery::build(app.issuer()))
}

pub async fn jwks(state_ext: AppStateExt) -> Result<Response, HttpError> {
    let axum::extract::State(app) = state_ext;
    let body = jwks::build(&app.db).map_err(|e| HttpError::api(CoreError::from(e)))?;
    let mut resp = Json(body).into_response();
    resp.headers_mut()
        .insert(header::CACHE_CONTROL, "public, max-age=300".parse().unwrap_or_else(|_| "no-store".parse().expect("static")));
    Ok(resp)
}

// ---------- /authorize ----------

#[derive(Debug, Deserialize)]
pub struct AuthorizeQuery {
    pub client_id: String,
    pub redirect_uri: String,
    pub response_type: String,
    pub scope: Option<String>,
    pub state: Option<String>,
    pub nonce: Option<String>,
    pub code_challenge: String,
    pub code_challenge_method: String,
}

/// `GET /oauth2/authorize`. If the user is not authenticated, redirect them
/// to the login page with a `next` parameter pointing back here. Otherwise
/// validate the request and immediately issue an authorization code.
///
/// (sui-id deliberately does not show a separate "consent" screen in the
/// minimal version — see the spec's §11.5/§13.2 commentary.)
pub async fn authorize(
    state_ext: AppStateExt,
    Query(q): Query<AuthorizeQuery>,
    headers: HeaderMap,
) -> Result<Response, HttpError> {
    let axum::extract::State(app) = state_ext;

    // Resolve a session if one exists; otherwise redirect to login.
    let (session_user, session_auth_methods) = match resolve_session(&app, &headers).await {
        Some(state) => state,
        None => {
            let qs = build_query_string(&q);
            let next = format!("/oauth2/authorize?{qs}");
            let encoded = utf8_percent_encode(&next, NON_ALPHANUMERIC).to_string();
            return Ok(Redirect::to(&format!("/admin/login?next={encoded}")).into_response());
        }
    };

    let client_id = ClientId::from_str(&q.client_id).map_err(|_| {
        HttpError::html(CoreError::Protocol {
            code: ProtocolError::InvalidClient,
            description: "client_id is not a valid identifier".into(),
        })
    })?;

    let scope = q.scope.clone().unwrap_or_else(|| "openid".into());
    let params = AuthorizeParams {
        client_id,
        redirect_uri: q.redirect_uri.clone(),
        response_type: q.response_type.clone(),
        scope,
        state: q.state.clone(),
        nonce: q.nonce.clone(),
        code_challenge: q.code_challenge.clone(),
        code_challenge_method: q.code_challenge_method.clone(),
    };

    let accepted = authorize::begin_authorization(&app.db, params).map_err(HttpError::html)?;
    let redirect = authorize::complete_authorization(
        &app.db,
        &app.clock,
        session_user,
        &session_auth_methods,
        accepted,
    )
    .map_err(HttpError::html)?;

    let mut url = redirect.redirect_uri;
    let sep = if url.contains('?') { '&' } else { '?' };
    url.push(sep);
    url.push_str("code=");
    url.push_str(&utf8_percent_encode(&redirect.code, NON_ALPHANUMERIC).to_string());
    if let Some(s) = redirect.state {
        url.push_str("&state=");
        url.push_str(&utf8_percent_encode(&s, NON_ALPHANUMERIC).to_string());
    }
    Ok(Redirect::to(&url).into_response())
}

async fn resolve_session(
    app: &crate::AppState,
    headers: &HeaderMap,
) -> Option<(sui_id_shared::ids::UserId, Vec<sui_id_shared::AuthMethod>)> {
    use axum_extra::extract::cookie::CookieJar;
    use sui_id_shared::ids::SessionId;
    let jar = CookieJar::from_headers(headers);
    let raw = jar.get(crate::handlers::SESSION_COOKIE)?.value().to_string();
    let id = SessionId::from_str(&raw).ok()?;
    // Resolve hits the database both for validity (expiry / revocation)
    // and to fetch the recorded auth_methods. Cheaper to do it once
    // here than to call resolve and then re-fetch the row.
    let session = sui_id_store::repos::sessions::get(&app.db, id).ok()?;
    if session.revoked_at.is_some() || session.expires_at < app.clock.now() {
        return None;
    }
    Some((session.user_id, session.auth_methods))
}

fn build_query_string(q: &AuthorizeQuery) -> String {
    let mut parts: Vec<(&str, String)> = vec![
        ("client_id", q.client_id.clone()),
        ("redirect_uri", q.redirect_uri.clone()),
        ("response_type", q.response_type.clone()),
        ("code_challenge", q.code_challenge.clone()),
        ("code_challenge_method", q.code_challenge_method.clone()),
    ];
    if let Some(s) = &q.scope {
        parts.push(("scope", s.clone()));
    }
    if let Some(s) = &q.state {
        parts.push(("state", s.clone()));
    }
    if let Some(n) = &q.nonce {
        parts.push(("nonce", n.clone()));
    }
    parts
        .into_iter()
        .map(|(k, v)| format!("{k}={}", utf8_percent_encode(&v, NON_ALPHANUMERIC)))
        .collect::<Vec<_>>()
        .join("&")
}

// ---------- /token ----------

#[derive(Debug, Deserialize)]
pub struct TokenForm {
    pub grant_type: String,
    pub code: Option<String>,
    pub redirect_uri: Option<String>,
    pub client_id: Option<String>,
    pub client_secret: Option<String>,
    pub code_verifier: Option<String>,
    pub refresh_token: Option<String>,
}

#[derive(Debug, Serialize)]
pub struct TokenResponse {
    pub access_token: String,
    pub token_type: &'static str,
    pub expires_in: i64,
    pub refresh_token: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub id_token: Option<String>,
    pub scope: Option<String>,
}

pub async fn token(
    state_ext: AppStateExt,
    crate::handlers::ClientIp(ip): crate::handlers::ClientIp,
    headers: HeaderMap,
    Form(form): Form<TokenForm>,
) -> Result<Response, HttpError> {
    let axum::extract::State(app) = state_ext;
    crate::handlers::enforce_rate_limit(
        &app.limiters,
        &app.clock,
        crate::handlers::RateLimitKey::Token,
        ip,
        crate::handlers::ErrorAs::Json,
    )?;

    let basic = parse_basic_auth(&headers);
    let (header_client, header_secret) = match basic {
        Some((c, s)) => (Some(c), Some(s)),
        None => (None, None),
    };
    let client_id_raw = header_client
        .or_else(|| form.client_id.clone())
        .ok_or_else(|| {
            HttpError::api(CoreError::Protocol {
                code: ProtocolError::InvalidClient,
                description: "client_id is required".into(),
            })
        })?;
    let client_id = ClientId::from_str(&client_id_raw).map_err(|_| {
        HttpError::api(CoreError::Protocol {
            code: ProtocolError::InvalidClient,
            description: "client_id is not a valid identifier".into(),
        })
    })?;
    let client_secret = header_secret.or_else(|| form.client_secret.clone());

    let lifetimes = app.token_lifetimes();
    let issuer = app.issuer().to_owned();
    let ctx = IssuanceContext {
        issuer: &issuer,
        lifetimes,
    };

    let set = match form.grant_type.as_str() {
        "authorization_code" => {
            let code = form.code.ok_or_else(|| {
                HttpError::api(CoreError::Protocol {
                    code: ProtocolError::InvalidRequest,
                    description: "code is required".into(),
                })
            })?;
            let redirect_uri = form.redirect_uri.ok_or_else(|| {
                HttpError::api(CoreError::Protocol {
                    code: ProtocolError::InvalidRequest,
                    description: "redirect_uri is required".into(),
                })
            })?;
            let code_verifier = form.code_verifier.ok_or_else(|| {
                HttpError::api(CoreError::Protocol {
                    code: ProtocolError::InvalidRequest,
                    description: "code_verifier is required (PKCE)".into(),
                })
            })?;
            authorize::exchange_code(
                &app.db,
                &app.clock,
                ctx,
                CodeExchangeRequest {
                    code,
                    redirect_uri,
                    client_id,
                    client_secret,
                    code_verifier,
                },
            )
            .map_err(HttpError::api)?
        }
        "refresh_token" => {
            let refresh_token = form.refresh_token.ok_or_else(|| {
                HttpError::api(CoreError::Protocol {
                    code: ProtocolError::InvalidRequest,
                    description: "refresh_token is required".into(),
                })
            })?;
            authorize::exchange_refresh(
                &app.db,
                &app.clock,
                ctx,
                RefreshExchangeRequest {
                    refresh_token,
                    client_id,
                    client_secret,
                },
            )
            .map_err(HttpError::api)?
        }
        other => {
            return Err(HttpError::api(CoreError::Protocol {
                code: ProtocolError::UnsupportedGrantType,
                description: format!("unsupported grant_type: {other}"),
            }));
        }
    };

    let resp = TokenResponse {
        access_token: set.access_token,
        token_type: "Bearer",
        expires_in: set.access_expires_in,
        refresh_token: set.refresh_token,
        id_token: set.id_token,
        scope: None,
    };
    let mut out = Json(resp).into_response();
    out.headers_mut().insert(
        header::CACHE_CONTROL,
        "no-store".parse().expect("static header value"),
    );
    out.headers_mut().insert(
        header::PRAGMA,
        "no-cache".parse().expect("static header value"),
    );
    Ok(out)
}

pub(crate) fn parse_basic_auth(headers: &HeaderMap) -> Option<(String, String)> {
    let raw = headers.get(header::AUTHORIZATION)?.to_str().ok()?;
    let body = raw.strip_prefix("Basic ")?;
    let mut buf = vec![0u8; body.len()];
    let n = Base64::decode(body, &mut buf).ok()?.len();
    buf.truncate(n);
    let s = String::from_utf8(buf).ok()?;
    let (id, secret) = s.split_once(':')?;
    Some((
        percent_decode(id).unwrap_or_else(|| id.to_owned()),
        percent_decode(secret).unwrap_or_else(|| secret.to_owned()),
    ))
}

fn percent_decode(s: &str) -> Option<String> {
    percent_encoding::percent_decode_str(s)
        .decode_utf8()
        .ok()
        .map(|c| c.into_owned())
}

// ---------- /userinfo ----------

#[derive(Debug, Serialize)]
pub struct UserInfo {
    pub sub: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub preferred_username: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
}

/// Standard OIDC userinfo endpoint. Authenticates via the Bearer access
/// token: we verify the JWT signature against our own JWKS, check expiry,
/// and look up the subject.
pub async fn userinfo(
    state_ext: AppStateExt,
    headers: HeaderMap,
) -> Result<Response, HttpError> {
    let axum::extract::State(app) = state_ext;
    let raw = headers
        .get(header::AUTHORIZATION)
        .and_then(|v| v.to_str().ok())
        .and_then(|s| s.strip_prefix("Bearer "))
        .ok_or_else(|| HttpError::api(CoreError::Unauthenticated))?;

    let claims = sui_id_core::tokens::verify_access_token(&app.db, &app.clock, raw)
        .map_err(HttpError::api)?;
    // RFC 7009: a revoked access token must stop being honoured at
    // protected endpoints. We consult the deny-list before serving.
    if sui_id_store::repos::revoked_access_tokens::is_revoked(&app.db, &claims.jti)
        .map_err(|e| HttpError::api(CoreError::from(e)))?
    {
        return Err(HttpError::api(CoreError::Unauthenticated));
    }
    let uid: sui_id_shared::ids::UserId = claims
        .sub
        .parse()
        .map_err(|_| HttpError::api(CoreError::Internal))?;
    let row = users::get(&app.db, uid).map_err(|_| HttpError::api(CoreError::Unauthenticated))?;
    if row.is_disabled || row.is_deleted {
        return Err(HttpError::api(CoreError::Unauthenticated));
    }
    // OIDC Core §5.3.2 SHOULD: userinfo is per-user data and must
    // not be cached by intermediaries. Without `no-store` a CDN or
    // shared proxy could serve one user's claims to another.
    let mut resp = Json(UserInfo {
        sub: row.id.to_string(),
        preferred_username: Some(row.username),
        name: row.display_name,
    })
    .into_response();
    resp.headers_mut().insert(
        header::CACHE_CONTROL,
        "no-store".parse().expect("static header value"),
    );
    Ok(resp)
}

// ---------- /logout (RP-initiated) ----------

/// Query parameters for the RP-initiated logout endpoint, per
/// OpenID Connect RP-Initiated Logout 1.0.
#[derive(Debug, Deserialize)]
pub struct LogoutQuery {
    /// Optional ID Token previously issued to this RP. Used as a hint for
    /// which user to log out and verified against our JWKS.
    pub id_token_hint: Option<String>,
    /// Where to send the user after logout. Must match a `redirect_uris`
    /// entry on the client referenced by `id_token_hint`. Without a hint
    /// we cannot validate against a client and must reject the parameter.
    pub post_logout_redirect_uri: Option<String>,
    /// Opaque value echoed back to the RP for CSRF protection.
    pub state: Option<String>,
    /// Optional `client_id`. Some RPs send it explicitly; we accept it as
    /// a fallback when there is no `id_token_hint`.
    pub client_id: Option<String>,
}

/// `GET /oauth2/logout`. End the user's session and, if a valid
/// `post_logout_redirect_uri` was supplied, redirect to it.
pub async fn logout(
    state_ext: AppStateExt,
    Query(q): Query<LogoutQuery>,
    headers: HeaderMap,
) -> Result<Response, HttpError> {
    let axum::extract::State(app) = state_ext;
    let mut user_id: Option<sui_id_shared::ids::UserId> = None;
    let mut hinted_client: Option<sui_id_shared::ids::ClientId> = None;

    // Prefer the id_token_hint when present.
    if let Some(token) = &q.id_token_hint {
        // Verify the signature; accept expired tokens so the user can still
        // log out after their session has gone stale.
        if let Ok(claims) = sui_id_core::tokens::verify_id_token(&app.db, &app.clock, token, true) {
            user_id = claims.sub.parse().ok();
            hinted_client = claims.aud.parse().ok();
        } else {
            tracing::warn!("logout: id_token_hint failed signature verification");
        }
    }

    // Fall back to the session cookie for the user identification.
    if user_id.is_none() {
        if let Some(cookie) = headers
            .get(header::COOKIE)
            .and_then(|v| v.to_str().ok())
        {
            for part in cookie.split(';') {
                let part = part.trim();
                if let Some(value) = part.strip_prefix("sui_id_session=") {
                    if let Ok(sid) = value.parse() {
                        if let Ok(uid) = sui_id_core::session::resolve(&app.db, &app.clock, sid) {
                            user_id = Some(uid);
                        }
                    }
                }
            }
        }
    }

    if let Some(uid) = user_id {
        sui_id_core::session::logout_user(&app.db, &app.clock, uid)
            .map_err(HttpError::html)?;
    }

    // Validate the post_logout_redirect_uri against the hinted client.
    //
    // Resolution order:
    //   1. Match against the client's own `post_logout_redirect_uris`
    //      (added in v0.6.0 via migration 0002). This is the standards-
    //      blessed list: an RP that registered logout URIs explicitly
    //      should only receive logout redirects to those URIs.
    //   2. If the client has *no* logout URIs registered, fall back to
    //      its `redirect_uris` for backwards compatibility with clients
    //      that pre-date this feature. We log a deprecation note when
    //      the fallback is taken.
    let redirect_target = match (
        q.post_logout_redirect_uri.as_deref(),
        hinted_client.or_else(|| q.client_id.as_deref().and_then(|s| s.parse().ok())),
    ) {
        (Some(uri), Some(cid)) => match sui_id_store::repos::clients::get(&app.db, cid) {
            Ok(client) if !client.is_disabled && !client.is_deleted => {
                if !client.post_logout_redirect_uris.is_empty() {
                    if client.post_logout_redirect_uris.iter().any(|u| u == uri) {
                        Some(uri.to_owned())
                    } else {
                        None
                    }
                } else if client.redirect_uris.iter().any(|u| u == uri) {
                    tracing::warn!(
                        client_id = %cid,
                        "logout: falling back to redirect_uris for post_logout_redirect_uri \
                         match (deprecated; register post_logout_redirect_uris on the client)"
                    );
                    Some(uri.to_owned())
                } else {
                    None
                }
            }
            _ => None,
        },
        _ => None,
    };

    // Always clear the session cookie before redirecting / responding.
    let clear_cookie = format!(
        "sui_id_session=; Path=/; HttpOnly; SameSite=Lax; Max-Age=0{}",
        if app.config.server.cookie_secure { "; Secure" } else { "" }
    );

    if let Some(mut url) = redirect_target {
        if let Some(state) = q.state.as_deref() {
            let sep = if url.contains('?') { '&' } else { '?' };
            url.push(sep);
            url.push_str("state=");
            url.push_str(&utf8_percent_encode(state, NON_ALPHANUMERIC).to_string());
        }
        let mut resp = Redirect::to(&url).into_response();
        if let Ok(v) = axum::http::HeaderValue::from_str(&clear_cookie) {
            resp.headers_mut().insert(header::SET_COOKIE, v);
        }
        return Ok(resp);
    }

    // No (or invalid) post_logout_redirect_uri: render a small confirmation page.
    let body = "<!DOCTYPE html><html><head><meta charset=\"utf-8\"><title>Signed out</title></head>\
                <body><h1>Signed out</h1><p>You have been signed out of sui-id.</p>\
                <p><a href=\"/admin/login\">Sign in again</a></p></body></html>";
    let mut resp = (axum::http::StatusCode::OK, axum::response::Html(body)).into_response();
    if let Ok(v) = axum::http::HeaderValue::from_str(&clear_cookie) {
        resp.headers_mut().insert(header::SET_COOKIE, v);
    }
    Ok(resp)
}
