/*
 * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
 *
 * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
 */

use super::{DeviceAuthResponse, FormData, MAX_POST_LEN, OAuthCode, PkceCodeChallenge};
use crate::auth::oauth::{OAuthStatus, openid::OpenIdHandler};
use common::{
    KV_OAUTH, Server,
    auth::{
        AuthRequest,
        oauth::{CLIENT_ID_MAX_LEN, DEVICE_CODE_LEN, USER_CODE_ALPHABET, USER_CODE_LEN},
    },
};
use directory::Credentials;
use http_proto::*;
use std::future::Future;
use store::{
    Serialize,
    dispatch::lookup::KeyValue,
    write::{Archive, Archiver},
};
use store::{
    rand::{
        Rng,
        distr::{Alphanumeric, StandardUniform},
        rng,
    },
    write::AlignedBytes,
};
use trc::AddContext;
use utils::DomainPart;

#[derive(Debug, serde::Serialize)]
pub struct OAuthMetadata {
    pub issuer: String,
    pub token_endpoint: String,
    pub authorization_endpoint: String,
    pub device_authorization_endpoint: String,
    pub registration_endpoint: String,
    pub introspection_endpoint: String,
    pub grant_types_supported: &'static [&'static str],
    pub response_types_supported: &'static [&'static str],
    pub scopes_supported: &'static [&'static str],
    pub code_challenge_methods_supported: &'static [&'static str],
}

pub trait OAuthApiHandler: Sync + Send {
    fn handle_discover_request(
        &self,
        session: &HttpSessionData,
        account_name: &str,
    ) -> impl Future<Output = trc::Result<HttpResponse>> + Send;

    fn handle_login_request(
        &self,
        session: &HttpSessionData,
        body: Vec<u8>,
    ) -> impl Future<Output = trc::Result<HttpResponse>> + Send;

    fn handle_device_auth(
        &self,
        req: &mut HttpRequest,
        session: &HttpSessionData,
    ) -> impl Future<Output = trc::Result<HttpResponse>> + Send;

    fn handle_oauth_metadata(&self) -> impl Future<Output = trc::Result<HttpResponse>> + Send;
}

#[derive(Debug, serde::Serialize, serde::Deserialize)]
#[serde(tag = "type")]
#[serde(rename_all = "camelCase")]
pub enum LoginRequest {
    #[serde(rename_all = "camelCase")]
    AuthCode {
        account_name: String,
        account_secret: String,
        #[serde(skip_serializing_if = "Option::is_none")]
        #[serde(default)]
        mfa_token: Option<String>,
        client_id: String,
        #[serde(default)]
        redirect_uri: Option<String>,
        #[serde(default)]
        nonce: Option<String>,
        #[serde(default)]
        scope: Option<String>,
        #[serde(default)]
        code_challenge: Option<String>,
        #[serde(default)]
        code_challenge_method: Option<String>,
        #[serde(default)]
        state: Option<String>,
    },
    #[serde(rename_all = "camelCase")]
    AuthDevice {
        account_name: String,
        account_secret: String,
        #[serde(skip_serializing_if = "Option::is_none")]
        #[serde(default)]
        mfa_token: Option<String>,
        code: String,
    },
}

#[derive(Debug, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
#[serde(tag = "type")]
#[serde(rename_all = "camelCase")]
pub enum LoginResponse {
    Authenticated { client_code: String },
    Verified,
    MfaRequired,
    Failure,
}

impl OAuthApiHandler for Server {
    async fn handle_discover_request(
        &self,
        session: &HttpSessionData,
        account_name: &str,
    ) -> trc::Result<HttpResponse> {
        let account_name = account_name.trim().to_lowercase();
        if let Some(domain_name) = account_name.try_domain_part()
            && let Some(endpoint) = self
                .get_directory_for_domain(domain_name)
                .await?
                .and_then(|directory| directory.oidc_discovery_document())
        {
            Ok(JsonResponse::new(&endpoint.document)
                .no_cache()
                .into_http_response())
        } else {
            self.handle_oidc_metadata(!session.is_tls).await
        }
    }

    async fn handle_login_request(
        &self,
        session: &HttpSessionData,
        body: Vec<u8>,
    ) -> trc::Result<HttpResponse> {
        let request = serde_json::from_slice::<LoginRequest>(&body).map_err(|err| {
            trc::EventType::Resource(trc::ResourceEvent::BadParameters).from_json_error(err)
        })?;

        let response = match request {
            LoginRequest::AuthCode {
                account_name,
                account_secret,
                mfa_token,
                client_id,
                redirect_uri,
                nonce,
                code_challenge,
                code_challenge_method,
                ..
            } => {
                // Validate clientId
                if client_id.len() > CLIENT_ID_MAX_LEN {
                    return Err(trc::AuthEvent::Error
                        .into_err()
                        .details("Client ID is too long."));
                } else if redirect_uri
                    .as_ref()
                    .is_some_and(|uri| uri.starts_with("http://"))
                {
                    #[cfg(not(feature = "dev_mode"))]
                    if !self.registry().is_recovery_mode() && code_challenge.is_none() {
                        return Err(trc::AuthEvent::Error
                            .into_err()
                            .details("Redirect URI must be HTTPS."));
                    }
                }

                // Parse and validate PKCE challenge (RFC 7636).
                let pkce_challenge = match code_challenge {
                    Some(challenge) => {
                        // Default to "plain" when the method is omitted, per RFC 7636 4.3.
                        match code_challenge_method.as_deref().unwrap_or("plain") {
                            "S256" => PkceCodeChallenge::S256(challenge),
                            "plain" => PkceCodeChallenge::Plain(challenge),
                            _ => {
                                return Err(trc::AuthEvent::Error
                                    .into_err()
                                    .details("Unsupported PKCE code_challenge_method."));
                            }
                        }
                    }
                    None => PkceCodeChallenge::None,
                };

                // Authenticate
                match self
                    .authenticate(&AuthRequest {
                        credentials: Credentials::Basic {
                            username: account_name,
                            secret: account_secret,
                            mfa_token,
                        },
                        session_id: session.session_id,
                        remote_ip: session.remote_ip,
                    })
                    .await
                {
                    Ok(access_token) => {
                        // Generate client code
                        let client_code = rng()
                            .sample_iter(Alphanumeric)
                            .take(DEVICE_CODE_LEN)
                            .map(char::from)
                            .collect::<String>();

                        // Serialize OAuth code
                        let value = Archiver::new(OAuthCode {
                            status: OAuthStatus::Authorized,
                            account_id: access_token.account_id(),
                            client_id,
                            nonce,
                            params: redirect_uri.unwrap_or_default(),
                            code_challenge: pkce_challenge,
                        })
                        .untrusted()
                        .serialize()
                        .caused_by(trc::location!())?;

                        // Insert client code
                        self.in_memory_store()
                            .key_set(
                                KeyValue::with_prefix(KV_OAUTH, client_code.as_bytes(), value)
                                    .expires(self.core.oauth.oauth_expiry_auth_code),
                            )
                            .await?;

                        LoginResponse::Authenticated { client_code }
                    }
                    Err(err) => match *err.as_ref() {
                        trc::EventType::Auth(trc::AuthEvent::MfaRequired) => {
                            trc::error!(err.span_id(session.session_id));
                            LoginResponse::MfaRequired
                        }
                        trc::EventType::Auth(_) => {
                            trc::error!(err.span_id(session.session_id));
                            LoginResponse::Failure
                        }
                        trc::EventType::Security(_) => {
                            trc::error!(err.span_id(session.session_id));
                            LoginResponse::Failure
                        }
                        _ => {
                            return Err(err);
                        }
                    },
                }
            }
            LoginRequest::AuthDevice {
                account_name,
                account_secret,
                mfa_token,
                code,
            } => {
                // Obtain code
                let mut result = LoginResponse::Failure;
                if let Some(auth_code_) = self
                    .in_memory_store()
                    .key_get::<Archive<AlignedBytes>>(KeyValue::<()>::build_key(
                        KV_OAUTH,
                        code.as_bytes(),
                    ))
                    .await?
                {
                    let oauth = auth_code_
                        .unarchive::<OAuthCode>()
                        .caused_by(trc::location!())?;
                    if oauth.status == OAuthStatus::Pending {
                        // Authenticate
                        match self
                            .authenticate(&AuthRequest {
                                credentials: Credentials::Basic {
                                    username: account_name,
                                    secret: account_secret,
                                    mfa_token,
                                },
                                session_id: session.session_id,
                                remote_ip: session.remote_ip,
                            })
                            .await
                        {
                            Ok(access_token) => {
                                let new_oauth_code = OAuthCode {
                                    status: OAuthStatus::Authorized,
                                    account_id: access_token.account_id(),
                                    client_id: oauth.client_id.to_string(),
                                    nonce: oauth.nonce.as_ref().map(|s| s.to_string()),
                                    params: Default::default(),
                                    code_challenge: PkceCodeChallenge::None,
                                };

                                // Delete issued user code
                                self.in_memory_store()
                                    .key_delete(KeyValue::<()>::build_key(
                                        KV_OAUTH,
                                        code.as_bytes(),
                                    ))
                                    .await?;

                                // Update device code status
                                self.in_memory_store()
                                    .key_set(
                                        KeyValue::with_prefix(
                                            KV_OAUTH,
                                            oauth.params.as_bytes(),
                                            Archiver::new(new_oauth_code)
                                                .untrusted()
                                                .serialize()
                                                .caused_by(trc::location!())?,
                                        )
                                        .expires(self.core.oauth.oauth_expiry_auth_code),
                                    )
                                    .await?;

                                result = LoginResponse::Verified;
                            }
                            Err(err) => match *err.as_ref() {
                                trc::EventType::Auth(trc::AuthEvent::MfaRequired) => {
                                    trc::error!(err.span_id(session.session_id));
                                    result = LoginResponse::MfaRequired;
                                }
                                trc::EventType::Auth(_) => {
                                    trc::error!(err.span_id(session.session_id));
                                    result = LoginResponse::Failure;
                                }
                                trc::EventType::Security(_) => {
                                    trc::error!(err.span_id(session.session_id));
                                    result = LoginResponse::Failure;
                                }
                                _ => {
                                    return Err(err);
                                }
                            },
                        }
                    }
                }

                result
            }
        };

        Ok(JsonResponse::new(response).no_cache().into_http_response())
    }

    async fn handle_device_auth(
        &self,
        req: &mut HttpRequest,
        session: &HttpSessionData,
    ) -> trc::Result<HttpResponse> {
        // Parse form
        let mut form_data = FormData::from_request(req, MAX_POST_LEN, session.session_id).await?;
        let client_id = form_data
            .remove("client_id")
            .filter(|client_id| client_id.len() <= CLIENT_ID_MAX_LEN)
            .ok_or_else(|| {
                trc::ResourceEvent::BadParameters
                    .into_err()
                    .details("Client ID is missing.")
            })?;
        let nonce = form_data.remove("nonce");

        // Generate device code
        let device_code = rng()
            .sample_iter(Alphanumeric)
            .take(DEVICE_CODE_LEN)
            .map(char::from)
            .collect::<String>();

        // Generate user code
        let mut user_code = String::with_capacity(USER_CODE_LEN + 1);
        for (pos, ch) in rng()
            .sample_iter(StandardUniform)
            .take(USER_CODE_LEN)
            .map(|v: u64| char::from(USER_CODE_ALPHABET[v as usize % USER_CODE_ALPHABET.len()]))
            .enumerate()
        {
            if pos == USER_CODE_LEN / 2 {
                user_code.push('-');
            }
            user_code.push(ch);
        }

        // Add OAuth status
        let oauth_code = Archiver::new(OAuthCode {
            status: OAuthStatus::Pending,
            account_id: u32::MAX,
            client_id,
            nonce,
            params: device_code.clone(),
            code_challenge: PkceCodeChallenge::None,
        })
        .untrusted()
        .serialize()
        .caused_by(trc::location!())?;

        // Insert device code
        self.in_memory_store()
            .key_set(
                KeyValue::with_prefix(KV_OAUTH, device_code.as_bytes(), oauth_code.clone())
                    .expires(self.core.oauth.oauth_expiry_user_code),
            )
            .await?;

        // Insert user code
        self.in_memory_store()
            .key_set(
                KeyValue::with_prefix(KV_OAUTH, user_code.as_bytes(), oauth_code)
                    .expires(self.core.oauth.oauth_expiry_user_code),
            )
            .await?;

        // Build response
        let base_url = &self.core.network.http.url_https;
        Ok(JsonResponse::new(DeviceAuthResponse {
            verification_uri: format!("{base_url}/device"),
            verification_uri_complete: format!("{base_url}/device/?code={user_code}"),
            device_code,
            user_code,
            expires_in: self.core.oauth.oauth_expiry_user_code,
            interval: 5,
        })
        .no_cache()
        .into_http_response())
    }

    async fn handle_oauth_metadata(&self) -> trc::Result<HttpResponse> {
        let base_url = &self.core.network.http.url_https;

        Ok(JsonResponse::new(OAuthMetadata {
            authorization_endpoint: format!("{base_url}/login",),
            token_endpoint: format!("{base_url}/auth/token"),
            device_authorization_endpoint: format!("{base_url}/auth/device"),
            introspection_endpoint: format!("{base_url}/auth/introspect"),
            registration_endpoint: format!("{base_url}/auth/register"),
            grant_types_supported: &[
                "authorization_code",
                "implicit",
                "urn:ietf:params:oauth:grant-type:device_code",
            ],
            response_types_supported: &["code", "id_token", "code token", "id_token token"],
            scopes_supported: &[
                "openid",
                "offline_access",
                "urn:ietf:params:jmap:core",
                "urn:ietf:params:jmap:mail",
                "urn:ietf:params:jmap:submission",
                "urn:ietf:params:jmap:vacationresponse",
            ],
            code_challenge_methods_supported: &["S256"],
            issuer: base_url.to_string(),
        })
        .into_http_response())
    }
}