use std::sync::Arc;

use axum::Extension;
use axum::extract::{Query, State};
use axum::response::Response;
use serde::Deserialize;

use crate::errors::ApiError;
use crate::oauth::{exchange_code, project, redirect_with_cookies, state};
use crate::services::AuthService;

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

pub async fn callback(
    State(service): State<AuthService>,
    Extension(verifier): Extension<Arc<jwt::JwtVerifier>>,
    Query(query): Query<CallbackQuery>,
) -> Result<Response, ApiError> {
    if let Some(error) = query.error {
        let description = query.error_description.unwrap_or_default();
        return Err(ApiError::bad_request(format!(
            "authorization failed: {error} {description}"
        )));
    }
    let Some(code) = query.code else {
        return Err(ApiError::bad_request("missing authorization code"));
    };
    let Some(state_token) = query.state else {
        return Err(ApiError::bad_request("missing authorization state"));
    };
    let Some(login_state) = state::take(&service, &state_token)
        .await
        .map_err(ApiError::from)?
    else {
        return Err(ApiError::bad_request("authorization state is unknown or expired"));
    };

    let redirect_uri = login_state
        .redirect_uri
        .clone()
        .unwrap_or_else(|| service.config().oauth.redirect_uri.clone());
    let tokens = exchange_code(&service, &code, &login_state.verifier, &redirect_uri)
        .await
        .map_err(ApiError::from)?;
    let id_token = tokens
        .id_token
        .as_deref()
        .ok_or_else(|| ApiError::bad_request("token response has no id_token"))?;
    let id_claims = verifier
        .verify_id(id_token)
        .await
        .map_err(|_| ApiError::from(crate::errors::AuthError::Unauthorized))?;
    if id_claims.nonce.as_deref() != Some(login_state.nonce.as_str()) {
        return Err(ApiError::bad_request("id_token nonce does not match"));
    }
    if let Some(at_hash) = id_claims.at_hash.as_deref()
        && at_hash != jwt::at_hash(&tokens.access_token)
    {
        return Err(ApiError::bad_request("id_token at_hash does not match"));
    }
    let claims = verifier
        .verify_access(&tokens.access_token)
        .await
        .map_err(|_| ApiError::from(crate::errors::AuthError::Unauthorized))?;
    if id_claims.sub != claims.sub.to_string() {
        return Err(ApiError::bad_request("id_token subject does not match"));
    }
    let kenauth = service
        .state
        .kenauth
        .as_ref()
        .ok_or_else(|| ApiError::internal("kenoauth client is not configured"))?;
    let public = kenauth
        .get_user_public(claims.sub)
        .await
        .map_err(|error| ApiError::from(crate::errors::AuthError::Internal(error.to_string())))?;
    project::upsert(&service, &public)
        .await
        .map_err(ApiError::from)?;

    let config = service.config();
    let mut cookies = vec![jwt::session_cookie(
        &config.oauth,
        &tokens.access_token,
        tokens.expires_in.max(60),
    )];
    if let Some(refresh) = tokens.refresh_token {
        cookies.push(jwt::refresh_cookie(&config.oauth, &refresh));
    }
    Ok(redirect_with_cookies(&login_state.redirect, &cookies))
}