use std::collections::HashMap;

use chrono::{DateTime, Duration, Local, Utc};
use rocket::serde::json::Json;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sha2::Digest;
use subtle::ConstantTimeEq;

use crate::api::auth::WithVerifiedFxaLogin;
use crate::api::{Empty, EMPTY};
use crate::db::DbConn;
use crate::types::oauth::{Scope, ScopeSet};
use crate::{
    api::{auth, serialize_dt},
    auth::Authenticated,
    crypto::{SecretBytes, SessionCredentials},
    types::{
        HawkKey, OauthAccessToken, OauthAccessType, OauthAuthorization, OauthAuthorizationID,
        OauthRefreshToken, OauthToken, OauthTokenID, SessionID, UserID,
    },
};

// MISSING get /oauth/client/{client_id}

pub(crate) struct OauthClient {
    pub(crate) id: &'static str,
    // NOTE not read so far, but good to have
    #[allow(dead_code)]
    pub(crate) name: &'static str,
    pub(crate) scopes: &'static [Scope<'static>],
}

const SESSION_SCOPE: Scope = Scope::borrowed("https://identity.mozilla.com/tokens/session");

// NOTE the telemetry scopes don't seem to be needed. since we'd have to give
// out keys for them (fxa does) we'll exclude them entirely.
// see fxa-auth-server/config/dev.json for lists of predefined clients and permissions.
pub(crate) const OAUTH_CLIENTS: [OauthClient; 2] = [
    OauthClient {
        id: "5882386c6d801776",
        name: "Firefox",
        scopes: &[
            Scope::borrowed("profile:write"),
            Scope::borrowed("https://identity.mozilla.com/apps/oldsync"),
            Scope::borrowed("https://identity.mozilla.com/tokens/session"),
            // "https://identity.mozilla.com/ids/ecosystem_telemetry",
        ],
    },
    OauthClient {
        id: "a2270f727f45f648",
        name: "Fenix",
        scopes: &[
            Scope::borrowed("profile"),
            Scope::borrowed("https://identity.mozilla.com/apps/oldsync"),
            Scope::borrowed("https://identity.mozilla.com/tokens/session"),
            // "https://identity.mozilla.com/ids/ecosystem_telemetry",
        ],
    },
];

// NOTE fxa dev config allows scoped keys only for:
//  - https://identity.mozilla.com/apps/notes
//  - https://identity.mozilla.com/apps/oldsync
//  - https://identity.mozilla.com/ids/ecosystem_telemetry
//  - https://identity.mozilla.com/apps/send
// we only implement sync because notes and send are dead and
// telemetry is of no use to us
const SCOPES_WITH_KEYS: [Scope; 1] = [Scope::borrowed("https://identity.mozilla.com/apps/oldsync")];

fn check_client_and_scopes(
    client_id: &str,
    scope: &ScopeSet,
) -> Result<&'static OauthClient, auth::Error> {
    let desc = match OAUTH_CLIENTS.iter().find(|&s| s.id == client_id) {
        Some(d) => d,
        None => return Err(auth::Error::UnknownClientID),
    };
    if !scope.is_allowed_by(desc.scopes) {
        return Err(auth::Error::ScopesNotAllowed);
    }
    Ok(desc)
}

#[derive(Debug, Deserialize)]
pub(crate) enum PkceChallengeType {
    S256,
}

#[derive(Debug, Deserialize)]
pub(crate) enum AuthResponseType {
    #[serde(rename = "code")]
    Code,
}

#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct OauthAuthReq {
    client_id: String,
    state: String,
    keys_jwe: Option<String>,
    scope: ScopeSet,
    access_type: OauthAccessType,
    // NOTE we don't support confidential clients, so PKCE is mandatory
    code_challenge: String,

    // MISSING redirect_uri
    // MISSING acr_value

    // for validation during deserialization only
    #[allow(dead_code)]
    code_challenge_method: PkceChallengeType,
    #[allow(dead_code)]
    response_type: AuthResponseType,
}

#[derive(Debug, Serialize)]
pub(crate) struct OauthAuthResp {
    code: OauthAuthorizationID,
    state: String,
    // MISSING redirect
}

#[post("/oauth/authorization", data = "<req>")]
pub(crate) async fn authorization(
    db: &DbConn,
    req: Authenticated<OauthAuthReq, WithVerifiedFxaLogin>,
) -> auth::Result<OauthAuthResp> {
    check_client_and_scopes(&req.body.client_id, &req.body.scope)?;
    let id = OauthAuthorizationID::random();
    db.add_oauth_authorization(
        &id,
        OauthAuthorization {
            user_id: req.context.uid,
            client_id: req.body.client_id,
            scope: req.body.scope,
            access_type: req.body.access_type,
            code_challenge: req.body.code_challenge,
            keys_jwe: req.body.keys_jwe,
            auth_at: req.context.created_at,
        },
    )
    .await?;
    Ok(Json(OauthAuthResp { code: id, state: req.body.state }))
}

#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct ScopedKeysReq {
    client_id: String,
    scope: ScopeSet,
}

#[derive(Debug, Serialize)]
#[allow(non_snake_case)]
pub(crate) struct ScopedKey {
    identifier: String,
    keyRotationSecret: &'static str,
    keyRotationTimestamp: u64,
}

#[post("/account/scoped-key-data", data = "<data>")]
pub(crate) async fn scoped_key_data(
    data: Authenticated<ScopedKeysReq, WithVerifiedFxaLogin>,
) -> auth::Result<HashMap<String, ScopedKey>> {
    check_client_and_scopes(&data.body.client_id, &data.body.scope)?;
    // like fxa we'll stub out key rotation handling entirely and return the same constants.
    Ok(Json(
        data.body
            .scope
            .split()
            .filter(|s| SCOPES_WITH_KEYS.contains(s))
            .map(|scope| {
                (
                    scope.to_string(),
                    ScopedKey {
                        identifier: scope.to_string(),
                        keyRotationSecret:
                            "0000000000000000000000000000000000000000000000000000000000000000",
                        keyRotationTimestamp: 0,
                    },
                )
            })
            .collect(),
    ))
}

#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct OauthDestroy {
    #[allow(dead_code)]
    client_id: String,
    token: OauthToken,
}

#[post("/oauth/destroy", data = "<data>")]
pub(crate) async fn destroy(db: &DbConn, data: Json<OauthDestroy>) -> auth::Result<Empty> {
    // MISSING api spec allows an optional basic auth header, but what for?
    // TODO fxa also checks the authorization header if present, but firefox doesn't send it
    // NOTE fxa does a constant-time check for whether the provided client_id matches that
    // of the token being destroyed. since we only support a public list we can skip that,
    // which also lets us more easily deal with firefox's habit of destroying tokens that are
    // already expired
    if db.delete_oauth_token(&data.token.hash()).await? {
        Ok(EMPTY)
    } else {
        Err(auth::Error::InvalidParameter)
    }
}

#[derive(Debug, Deserialize)]
#[serde(tag = "grant_type")]
enum TokenReqDetails {
    // we can't use deny_unknown_fields when flatten is involved, and multiple
    // flattens in the same struct cause problems if one of them is greedy (like map).
    // flatten an extra map into every variant instead and check each of them.
    #[serde(rename = "authorization_code")]
    AuthCode {
        code: OauthAuthorizationID,
        code_verifier: String,
        // NOTE only useful with redirect flows, which we kinda don't support at all
        #[allow(dead_code)]
        redirect_uri: Option<String>,
        #[serde(flatten)]
        extra: HashMap<String, Value>,
    },
    #[serde(rename = "refresh_token")]
    RefreshToken {
        refresh_token: OauthToken,
        scope: ScopeSet,
        #[serde(flatten)]
        extra: HashMap<String, Value>,
    },
    #[serde(rename = "fxa-credentials")]
    FxaCreds {
        scope: ScopeSet,
        access_type: Option<OauthAccessType>,
        #[serde(flatten)]
        extra: HashMap<String, Value>,
    },
}

impl TokenReqDetails {
    fn extra_is_empty(&self) -> bool {
        match self {
            TokenReqDetails::AuthCode { extra, .. } => extra.is_empty(),
            TokenReqDetails::RefreshToken { extra, .. } => extra.is_empty(),
            TokenReqDetails::FxaCreds { extra, .. } => extra.is_empty(),
        }
    }
}

// TODO log errors in all the places

#[derive(Debug, Deserialize)]
pub(crate) struct TokenReq {
    client_id: String,
    ttl: Option<u32>,
    #[serde(flatten)]
    details: TokenReqDetails,
    // MISSING client_secret
    // MISSING redirect_uri
    // MISSING ttl
    // MISSING ppid_seed
    // MISSING resource
}

#[derive(Debug, Serialize)]
pub(crate) enum TokenType {
    #[serde(rename = "bearer")]
    Bearer,
}

#[derive(Debug, Serialize)]
pub(crate) struct TokenResp {
    access_token: OauthToken,
    #[serde(skip_serializing_if = "Option::is_none")]
    refresh_token: Option<OauthToken>,
    // MISSING id_token
    #[serde(skip_serializing_if = "Option::is_none")]
    session_token: Option<String>,
    scope: ScopeSet,
    token_type: TokenType,
    expires_in: u32,
    #[serde(serialize_with = "serialize_dt")]
    auth_at: DateTime<Utc>,
    #[serde(skip_serializing_if = "Option::is_none")]
    keys_jwe: Option<String>,
}

#[post("/oauth/token", data = "<req>", rank = 1)]
pub(crate) async fn token_authenticated(
    db: &DbConn,
    req: Authenticated<TokenReq, WithVerifiedFxaLogin>,
) -> auth::Result<TokenResp> {
    match &req.body.details {
        TokenReqDetails::FxaCreds { .. } => (),
        _ => return Err(auth::Error::InvalidParameter),
    }
    token_impl(
        db,
        Some(req.context.uid),
        Some(req.context.created_at),
        req.body,
        None,
        Some(req.session.clone()),
    )
    .await
}

#[post("/oauth/token", data = "<req>", rank = 2)]
pub(crate) async fn token_unauthenticated(
    db: &DbConn,
    req: Json<TokenReq>,
) -> auth::Result<TokenResp> {
    let (parent_refresh, auth_at) = match &req.details {
        TokenReqDetails::RefreshToken { refresh_token, .. } => {
            let session = db.use_session_from_refresh(&refresh_token.hash()).await?;
            (Some(refresh_token.hash()), Some(session.1.created_at))
        },
        TokenReqDetails::AuthCode { .. } => (None, None),
        _ => return Err(auth::Error::InvalidParameter),
    };
    token_impl(db, None, auth_at, req.into_inner(), parent_refresh, None).await
}

async fn token_impl(
    db: &DbConn,
    user_id: Option<UserID>,
    auth_at: Option<DateTime<Utc>>,
    req: TokenReq,
    parent_refresh: Option<OauthTokenID>,
    parent_session: Option<SessionID>,
) -> auth::Result<TokenResp> {
    if !req.details.extra_is_empty() {
        return Err(auth::Error::InvalidParameter);
    }
    let ttl = req.ttl.unwrap_or(3600).clamp(0, 7 * 86400);

    let (auth_at, scope, keys_jwe, user_id, access_type) = match req.details {
        TokenReqDetails::AuthCode { code, code_verifier, .. } => {
            let auth = match db.take_oauth_authorization(&code).await {
                Ok(a) => a,
                Err(_) => return Err(auth::Error::InvalidAuthToken),
            };
            if !bool::from(auth.client_id.as_bytes().ct_eq(req.client_id.as_bytes())) {
                return Err(auth::Error::UnknownClientID);
            }
            let mut sha = sha2::Sha256::new();
            sha.update(code_verifier.as_bytes());
            let challenge = base64::encode_config(&sha.finalize(), base64::URL_SAFE_NO_PAD);
            if !bool::from(challenge.as_bytes().ct_eq(auth.code_challenge.as_bytes())) {
                return Err(auth::Error::InvalidParameter);
            }
            (auth.auth_at, auth.scope, auth.keys_jwe, auth.user_id, Some(auth.access_type))
        },
        TokenReqDetails::RefreshToken { refresh_token, scope, .. } => {
            let auth_at =
                auth_at.expect("oauth token requests with refresh token must set auth_at");
            let base = db.get_refresh_token(&refresh_token.hash()).await?;
            if !bool::from(base.client_id.as_bytes().ct_eq(req.client_id.as_bytes())) {
                return Err(auth::Error::UnknownClientID);
            }
            check_client_and_scopes(&req.client_id, &scope)?;
            if !base.scope.implies_all(&scope) {
                return Err(auth::Error::ScopesNotAllowed);
            }
            (auth_at, scope, None, base.user_id, None)
        },
        TokenReqDetails::FxaCreds { scope, access_type, .. } => {
            let user_id = user_id.expect("oauth token requests with fxa must set user_id");
            let auth_at = auth_at.expect("oauth token requests with fxa must set auth_at");
            check_client_and_scopes(&req.client_id, &scope)?;
            (auth_at, scope, None, user_id, access_type)
        },
    };

    let access_token = OauthToken::random();
    db.add_access_token(
        &access_token.hash(),
        OauthAccessToken {
            user_id: user_id.clone(),
            client_id: req.client_id.clone(),
            scope: scope.clone(),
            parent_refresh,
            parent_session,
            expires_at: (Local::now() + Duration::seconds(ttl.into())).into(),
        },
    )
    .await?;

    let (refresh_token, session_token) = if access_type == Some(OauthAccessType::Offline) {
        let (session_token, session_id) = if scope.implies(&SESSION_SCOPE) {
            let session_token = SecretBytes::generate();
            let session = SessionCredentials::derive(&session_token);
            let session_id = SessionID(session.token_id.0);
            db.add_session(session_id.clone(), &user_id, HawkKey(session.req_hmac_key), true, None)
                .await?;
            (Some(session_token.0), Some(SessionID(session.token_id.0)))
        } else {
            (None, None)
        };

        let refresh_token = OauthToken::random();
        db.add_refresh_token(
            &refresh_token.hash(),
            OauthRefreshToken {
                user_id,
                client_id: req.client_id,
                scope: scope.remove(&SESSION_SCOPE),
                session_id,
            },
        )
        .await?;
        (Some(refresh_token), session_token)
    } else {
        (None, None)
    };

    Ok(Json(TokenResp {
        access_token,
        refresh_token,
        session_token: session_token.map(hex::encode),
        scope: scope.remove(&SESSION_SCOPE),
        token_type: TokenType::Bearer,
        expires_in: ttl,
        auth_at,
        keys_jwe,
    }))
}