use axum::{
    Json,
    extract::{Query, State},
    http::StatusCode,
    response::{IntoResponse, Response},
};
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::sync::LazyLock;
use tracing::{error, info, warn};
use tranquil_pds::api::error::ApiError;
use tranquil_pds::auth::extractor::{Auth, Permissive};
use tranquil_pds::state::AppState;
use tranquil_pds::types::Did;
use tranquil_types::Nsid;

static CREATE_ACCOUNT_NSID: LazyLock<Nsid> =
    LazyLock::new(|| "com.atproto.server.createAccount".parse().unwrap());

const HOUR_SECS: i64 = 3600;
const MINUTE_SECS: i64 = 60;

static PROTECTED_METHODS: LazyLock<HashSet<&'static str>> = LazyLock::new(|| {
    [
        "com.atproto.admin.sendEmail",
        "com.atproto.identity.requestPlcOperationSignature",
        "com.atproto.identity.signPlcOperation",
        "com.atproto.identity.updateHandle",
        "com.atproto.server.activateAccount",
        "com.atproto.server.confirmEmail",
        "com.atproto.server.createAppPassword",
        "com.atproto.server.deactivateAccount",
        "com.atproto.server.getAccountInviteCodes",
        "com.atproto.server.getSession",
        "com.atproto.server.listAppPasswords",
        "com.atproto.server.requestAccountDelete",
        "com.atproto.server.requestEmailConfirmation",
        "com.atproto.server.requestEmailUpdate",
        "com.atproto.server.revokeAppPassword",
        "com.atproto.server.updateEmail",
    ]
    .into_iter()
    .collect()
});

#[derive(Deserialize)]
pub struct GetServiceAuthParams {
    pub aud: String,
    pub lxm: Option<Nsid>,
    pub exp: Option<i64>,
}

#[derive(Serialize)]
pub struct GetServiceAuthOutput {
    pub token: String,
}

fn is_valid_did_fragment(fragment: &str) -> bool {
    let bytes = fragment.as_bytes();
    let mut index = 0;
    while index < bytes.len() {
        let byte = bytes[index];
        if byte.is_ascii_alphanumeric()
            || matches!(
                byte,
                b'-' | b'.'
                    | b'_'
                    | b'~'
                    | b'!'
                    | b'$'
                    | b'&'
                    | b'\''
                    | b'('
                    | b')'
                    | b'*'
                    | b'+'
                    | b','
                    | b';'
                    | b'='
                    | b':'
                    | b'@'
                    | b'/'
                    | b'?'
            )
        {
            index += 1;
        } else if byte == b'%'
            && bytes
                .get(index + 1..index + 3)
                .is_some_and(|encoded| encoded.iter().all(u8::is_ascii_hexdigit))
        {
            index += 3;
        } else {
            return false;
        }
    }
    !fragment.is_empty()
}

fn validate_service_audience(aud: &str) -> Result<(), ApiError> {
    let did = match aud.split_once('#') {
        Some((did, fragment)) => {
            if fragment.contains('#') || !is_valid_did_fragment(fragment) {
                return Err(ApiError::InvalidRequest(
                    "aud must be a valid atproto DID or did#serviceId reference".into(),
                ));
            }
            did
        }
        None => aud,
    };
    let did = Did::new(did).map_err(|_| {
        ApiError::InvalidRequest(
            "aud must be a valid atproto DID or did#serviceId reference".into(),
        )
    })?;
    if !did.is_plc() && !did.is_web() {
        return Err(ApiError::InvalidRequest(
            "aud must be a valid atproto DID or did#serviceId reference".into(),
        ));
    }
    Ok(())
}

pub async fn get_service_auth(
    State(state): State<AppState>,
    auth: Auth<Permissive>,
    Query(params): Query<GetServiceAuthParams>,
) -> Response {
    if let Err(error) = validate_service_audience(&params.aud) {
        return error.into_response();
    }
    info!(
        did = %&auth.did,
        is_oauth = auth.is_oauth(),
        aud = %params.aud,
        lxm = ?params.lxm,
        "getServiceAuth called"
    );

    let key_bytes = match &auth.key_bytes {
        Some(kb) => kb.clone(),
        None => {
            warn!(did = %&auth.did, "getServiceAuth: no key_bytes in auth, fetching from DB");
            match state.repos.user.get_user_info_by_did(&auth.did).await {
                Ok(Some(info)) => match info.key_bytes {
                    Some(key_bytes_enc) => {
                        match tranquil_pds::config::decrypt_key(
                            &key_bytes_enc,
                            info.encryption_version,
                        ) {
                            Ok(key) => key,
                            Err(e) => {
                                error!(error = ?e, "Failed to decrypt user key for service auth");
                                return ApiError::AuthenticationFailed(Some(
                                    "Failed to get signing key".into(),
                                ))
                                .into_response();
                            }
                        }
                    }
                    None => {
                        return ApiError::AuthenticationFailed(Some(
                            "User has no signing key".into(),
                        ))
                        .into_response();
                    }
                },
                Ok(None) => {
                    return ApiError::AuthenticationFailed(Some("User has no signing key".into()))
                        .into_response();
                }
                Err(e) => {
                    error!(error = ?e, "DB error fetching user key");
                    return ApiError::AuthenticationFailed(Some(
                        "Failed to get signing key".into(),
                    ))
                    .into_response();
                }
            }
        }
    };

    let lxm = params.lxm.as_ref();

    if let Some(method) = lxm {
        if let Err(e) = tranquil_pds::auth::scope_check::check_rpc_scope(
            &auth.auth_source,
            auth.scope.as_deref(),
            &params.aud,
            method,
        ) {
            return e.into_response();
        }
    } else if auth.is_oauth() {
        let permissions = auth.permissions();
        if !permissions.has_full_access() {
            return ApiError::InvalidRequest(
                "OAuth tokens with granular scopes must specify an lxm parameter".into(),
            )
            .into_response();
        }
    }

    if auth.status.is_takendown() && lxm != Some(&*CREATE_ACCOUNT_NSID) {
        return ApiError::InvalidToken(Some("Bad token scope".into())).into_response();
    }

    if let Some(method) = lxm
        && PROTECTED_METHODS.contains(&method.as_str())
    {
        return ApiError::InvalidRequest(format!(
            "cannot request a service auth token for the following protected method: {}",
            method
        ))
        .into_response();
    }

    if let Some(exp) = params.exp {
        let now = chrono::Utc::now().timestamp();
        let diff = exp - now;

        if diff < 0 {
            return ApiError::InvalidRequest("expiration is in past".into()).into_response();
        }

        if diff > HOUR_SECS {
            return ApiError::InvalidRequest(
                "cannot request a token with an expiration more than an hour in the future".into(),
            )
            .into_response();
        }

        if lxm.is_none() && diff > MINUTE_SECS {
            return ApiError::InvalidRequest(
                "cannot request a method-less token with an expiration more than a minute in the future".into(),
            )
            .into_response();
        }
    }

    let service_token =
        match tranquil_pds::auth::create_service_token(&auth.did, &params.aud, lxm, &key_bytes) {
            Ok(t) => t,
            Err(e) => {
                error!("Failed to create service token: {:?}", e);
                return ApiError::InternalError(None).into_response();
            }
        };
    (
        StatusCode::OK,
        Json(GetServiceAuthOutput {
            token: service_token,
        }),
    )
        .into_response()
}