use std::{borrow::Borrow, collections::HashMap, iter::Iterator, time::SystemTime};

use crate::{
    config::{
        sso::{Registration, Template},
        IdpConfig,
    },
    service::sso::{
        templates, LoginToken, RegistrationInfo, RegistrationToken, ValidationData,
        REGISTRATION_EXPIRATION_SECS, SESSION_EXPIRATION_SECS, SSO_AUTH_EXPIRATION_SECS,
        SSO_SESSION_COOKIE,
    },
    services, utils, Error, Result, Ruma,
};
use axum::{
    extract::RawQuery,
    response::{AppendHeaders, IntoResponse, Redirect},
    RequestExt,
};
use axum_extra::{
    headers::{self, HeaderMapExt},
    TypedHeader,
};
use http::header;
use mas_oidc_client::{
    requests::{
        authorization_code::{self, AuthorizationRequestData, AuthorizationValidationData},
        jose::{self, JwtVerificationData},
        userinfo,
    },
    types::{
        client_credentials::ClientCredentials,
        errors::ClientError,
        iana::jose::JsonWebSignatureAlg,
        requests::{AccessTokenResponse, AuthorizationResponse},
    },
};
use rand::{rngs::StdRng, SeedableRng};
use ruma::{
    api::client::{
        error::ErrorKind,
        session::{self, sso_login, sso_login_with_provider},
    },
    events::GlobalAccountDataEventType,
    push, OwnedMxcUri, UserId,
};
use serde_json::Number;
use tracing::error;
use url::Url;

pub const CALLBACK_PATH: &str = "_matrix/client/unstable/sso/callback";

/// # `GET /_matrix/client/v3/login/sso/redirect`
///
/// Redirect the user to the SSO interface.
/// TODO: this should be removed once Ruma supports trailing slashes.
pub async fn get_sso_redirect(
    body: Ruma<sso_login::v3::Request>,
) -> Result<sso_login::v3::Response> {
    let sso_login_with_provider::v3::Response { location, cookie } =
        get_sso_redirect_with_provider(
            Ruma {
                body: sso_login_with_provider::v3::Request::new(
                    Default::default(),
                    body.redirect_url.clone(),
                ),
                ..body
            }
            .into(),
        )
        .await?;

    Ok(sso_login::v3::Response { location, cookie })
}

/// # `GET /_matrix/client/v3/login/sso/redirect/{idpId}`
///
/// Redirects the user to the SSO interface.
pub async fn get_sso_redirect_with_provider(
    body: Ruma<sso_login_with_provider::v3::Request>,
) -> Result<sso_login_with_provider::v3::Response> {
    let idp_ids: Vec<&str> = services()
        .globals
        .config
        .idps
        .iter()
        .map(Borrow::borrow)
        .collect();

    let provider = match &*idp_ids {
        [] => {
            return Err(Error::BadRequest(
                ErrorKind::forbidden(),
                "Single Sign-On is disabled.",
            ));
        }
        [idp_id] => services().sso.get(idp_id).expect("we know it exists"),
        [_, ..] => services().sso.get(&body.idp_id).ok_or_else(|| {
            Error::BadRequest(ErrorKind::InvalidParam, "Unknown identity provider.")
        })?,
    };

    let redirect_url = body
        .redirect_url
        .parse::<Url>()
        .map_err(|_| Error::BadRequest(ErrorKind::InvalidParam, "Invalid redirect_url."))?;

    let mut callback = services()
        .globals
        .well_known_client()
        .parse::<Url>()
        .map_err(|_| Error::bad_config("Invalid well_known_client url."))?;
    callback.set_path(CALLBACK_PATH);

    let (auth_url, validation_data) = authorization_code::build_authorization_url(
        provider.metadata.authorization_endpoint().clone(),
        AuthorizationRequestData::new(
            provider.config.client_id.clone(),
            provider.config.scopes.clone(),
            redirect_url,
        ),
        &mut StdRng::from_entropy(),
    )
    .map_err(|_| Error::BadRequest(ErrorKind::Unknown, "Failed to build authorization_url."))?;

    let signed = services().globals.sign_claims(&ValidationData::new(
        provider.borrow().to_string(),
        validation_data,
    ));

    Ok(sso_login_with_provider::v3::Response {
        location: auth_url.to_string(),
        cookie: Some(
            utils::build_cookie(
                SSO_SESSION_COOKIE,
                &signed,
                "/_conduit/client/sso/callback",
                Some(SSO_AUTH_EXPIRATION_SECS),
            )
            .to_string(),
        ),
    })
}

/// # `GET /_conduit/client/sso/callback`
///
/// Validate the authorization response received from the identity provider.
/// On success, generate a login token, add it to `redirectUrl` as a query and perform the redirect.
/// If this is the first login, register the user, possibly interactively through a fallback page.
pub async fn get_sso_callback(req: axum::extract::Request) -> Result<axum::response::Response> {
    let query = req.uri().query().ok_or_else(|| {
        Error::BadRequest(ErrorKind::MissingParam, "Empty authorization callback.")
    })?;

    let AuthorizationResponse {
        code,
        access_token,
        token_type,
        id_token,
        expires_in,
    } = serde_html_form::from_str::<AuthorizationResponse>(query).map_err(|_| {
        serde_html_form::from_str::<ClientError>(query).unwrap_or_else(|_| {
            error!("Failed to deserialize authorization callback: {}", callback);

            Error::BadRequest(
                ErrorKind::Unknown,
                "Failed to deserialize authorization callback.",
            )
        })
    })?;

    let cookie = req
        .extract::<Option<TypedHeader<headers::Cookie>>>()
        .await
        .map_err(|_| Error::BadRequest(ErrorKind::InvalidParam, "Invalid session cookie."))?
        .ok_or_else(|_| Error::BadRequest(ErrorKind::MissingParam, "Missing session cookie."))?;

    let ValidationData {
        provider,
        inner: validation_data,
    } = services()
        .globals
        .validate_claims(
            cookie.get(SSO_SESSION_COOKIE).ok_or_else(|| {
                Error::BadRequest(ErrorKind::MissingParam, "Missing value for session cookie.")
            })?,
            None,
        )
        .map_err(|e| {
            Error::BadRequest(ErrorKind::InvalidParam, "Invalid value for session cookie.")
        })?;

    let provider = services().sso.get(&provider).ok_or_else(|e| {
        Error::BadRequest(
            ErrorKind::InvalidParam,
            "Unknown provider for session cookie.",
        )
    })?;

    let IdpConfig {
        client_id,
        client_secret,
        auth_method,
        ..
    } = provider.config.clone();

    let credentials = match &auth_method {
        "basic" => ClientCredentials::ClientSecretBasic {
            client_id,
            client_secret,
        },
        "post" => ClientCredentials::ClientSecretPost {
            client_id,
            client_secret,
        },
        _ => todo!(),
    };
    let (
        AccessTokenResponse {
            access_token,
            refresh_token,
            token_type,
            expires_in,
            scope,
            ..
        },
        Some(id_token),
    ) = authorization_code::access_token_with_authorization_code(
        services().sso.service(),
        method,
        provider.metadata.token_endpoint(),
        code,
        validation_data,
        jwt_verification_data,
        SystemTime::now().into(),
        &mut StdRng::from_entropy(),
    )
    .await
    .map_err(|e| Error::bad_config("Failed to fetch access token."))?
    else {
        unreachable!("ID token should never be empty")
    };

    // let userinfo = provider.fetch_userinfo(&access_token, &id_token).await?;

    let mut userinfo = HashMap::default();
    if let Some(endpoint) = &provider.metadata.userinfo_endpoint {
        let ref jwks = jose::fetch_jwks(services().sso.service(), provider.metadata.jwks_uri())
            .await
            .map_err(|e| Error::bad_config("Failed to fetch signing keys for token endpoint."))?;
        let jwt_verification_data = Some(JwtVerificationData {
            jwks,
            issuer: &provider.config.issuer,
            client_id: credentials.client_id(),
            signing_algorithm: &JsonWebSignatureAlg::Rs256,
        });

        userinfo = userinfo::fetch_userinfo(
            services().sso.service(),
            endpoint,
            &access_token,
            jwt_verification_data,
            &id_token,
        )
        .await
        .map_err(|e| Error::bad_config("Failed to fetch claims for userinfo endpoint."))?;
    };

    let (_, mut claims) = id_token.into_parts();

    let subject = claims.get("sub").ok_or_else(|| {
        error!("Unique \"sub\" claim is missing from ID token: {claims:?}");

        Error::bad_config("Unique \"sub\" claim is missing from ID token.")
    })?;

    let subject = &subject
        .as_str()
        .map(str::to_owned)
        .or_else(|| subject.as_number().map(Number::to_string))
        .expect("unique claim should be a string or number");

    let redirect_uri = &validation_data.redirect_uri;

    if let Some(user_id) = services()
        .sso
        .user_from_claim(&validation_data.provider_id, subject)?
    {
        let login_token = LoginToken::new(validation_data.provider_id.to_owned(), user_id);

        let redirect_uri = redirect_with_login_token(redirect_uri.to_owned(), &login_token);

        return Ok((
            AppendHeaders(vec![(
                header::SET_COOKIE,
                utils::reset_cookie("sso-session").to_string(),
            )]),
            Redirect::temporary(redirect_uri.as_str()),
        )
            .into_response());
    }

    match provider.config.registration {
        Registration::Disabled => {
            return Err(Error::BadRequest(
                ErrorKind::forbidden(),
                "Single Sign-On registration is disabled.",
            ))
        }
        Registration::Automated => todo!(),
        Registration::Interactive => {}
    };

    let Template {
        username,
        displayname,
        avatar_url,
        email,
    } = &provider.config.template;
    let registration_info =
        RegistrationInfo::new(&claims, username, displayname, avatar_url, email);

    let signed = services()
        .globals
        .sign_macaroon(&RegistrationToken::new(
            validation_data.provider_id.clone(),
            subject.to_owned(),
            redirect_uri.to_owned(),
            registration_info,
        ))
        .expect("signing macaroons always works");

    let cookie = utils::build_cookie(
        "sso-registration",
        &signed,
        "/_conduit/client/sso/register",
        REGISTRATION_EXPIRATION_SECS,
    );

    Ok((
        AppendHeaders(vec![
            (header::SET_COOKIE, cookie.to_string()),
            (
                header::SET_COOKIE,
                utils::reset_cookie("sso-session").to_string(),
            ),
        ]),
        Redirect::temporary("/_conduit/client/sso/register"),
    )
        .into_response())
}

/// # `GET /_conduit/client/sso/pick_idp`
pub async fn pick_idp(RawQuery(query): RawQuery) -> impl IntoResponse {
    let providers: Vec<_> = services()
        .globals
        .config
        .sso
        .iter()
        .map(|p| p.inner.to_owned())
        .collect();

    let body = maud::html! {
        header {
            h1 { "Log in to " (services().globals.server_name()) }
            p { "Choose an identity provider to continue" }
        }
        main {
            ul .providers {
                @for provider in providers {
                    li {
                        a href={ "/_matrix/client/v3/login/sso/redirect/" (provider.id) "?" (query.as_deref().unwrap_or_default()) } {
                            @if let Some(url) = provider.icon.as_deref().and_then(utils::mxc_to_http) {
                                img src=(url);
                            }
                        }
                        span {
                            (provider.name)
                        }
                    }
                }
            }
        }
    };

    (
        [(header::CONTENT_TYPE, "text/html; charset=utf-8")],
        maud::html! {
            (templates::base("Pick Identity Provider", body))

            (templates::footer())
        },
    )
}

/// # `GET /_conduit/client/sso/register`
///
/// Serve a registration form with defaults based on the retrieved claims.
/// This endpoint is only available when interactive registration is enabled.
pub async fn get_sso_registration(
    cookie: TypedHeader<headers::Cookie>,
) -> Result<axum::response::Response> {
    let token = cookie.get("sso-registration").ok_or_else(|| {
        Error::BadRequest(
            ErrorKind::MissingParam,
            "Missing registration token cookie.",
        )
    })?;

    let registration_token: RegistrationToken = services()
        .globals
        .validate_macaroon(token, None)
        .map_err(|_| {
        Error::BadRequest(
            ErrorKind::InvalidParam,
            "Invalid registration token cookie.",
        )
    })?;

    let provider = services()
        .sso
        .get(&registration_token.provider_id)
        .map(|p| p.config.inner.to_owned())?;
    let server_name = services().globals.server_name();

    let RegistrationInfo {
        username,
        displayname,
        avatar_url,
        email,
    } = registration_token.info;

    let additional_info = (&displayname, &avatar_url, &email) != (&None, &None, &None);

    fn detail(title: &str, body: maud::Markup) -> maud::Markup {
        maud::html! {
            label .detail for=(title) {
                div .check-row {
                    span .name { (title) } " "
                    span .use { "use" }
                    input #(title) type="checkbox" name={(title)"-checkbox"} value=(true) checked;
                }
                (body)
            }
        }
    }

    let body = maud::html! {
        header {
            h1 { "Complete your registration at " (server_name) }
            p { "Confirm your details to finish creating your account." }
        }
        main {
            form .form #form method="post" {
                div .username-div #username-div {
                    label for="username-input" { "Username (required)" }
                    div .prefix { "@" }
                    input .username-input type="text" name="username"
                        value=(username) autofocus autocorrect="off" autocapitalize="none";
                    div .postfix { ":" (server_name) }
                }
                output .username-output for="username-input" { }

                @if additional_info {
                    section .additional-info {
                        h2 {
                            @if let Some(icon) = provider.icon.as_deref().and_then(utils::mxc_to_http) {
                                img src=(icon.to_string());
                            }
                            "Optional data from " (provider.name)
                        }
                        @if let Some(avatar_url) = avatar_url.as_ref() {
                            (detail("avatar", maud::html!{
                                img .avatar src=(avatar_url);
                            }))
                        }
                        @if let Some(displayname) = displayname.as_ref() {
                            (detail("displayname", maud::html!{
                                p .value { (displayname) };
                            }))
                        }
                        @if let Some(email) = email.as_ref() {
                            (detail("email", maud::html!{
                                p .value { (email) };
                            }))
                        }
                    }
                }

                input type="submit" value="Submit" .primary-button {}
            }
        }
    };

    Ok((
        [(header::CONTENT_TYPE, "text/html; charset=utf-8")],
        maud::html! {
            (templates::base("Register Account", body))

            (templates::footer())
        },
    )
        .into_response())
}

/// # `POST /_conduit/client/sso/register`
///
/// Submit the registration form.
pub async fn submit_sso_registration(
    cookie: TypedHeader<headers::Cookie>,
    axum::extract::Form(registration_info): axum::extract::Form<RegistrationInfo>,
) -> Result<axum::response::Response> {
    let token = cookie.get("sso-registration").ok_or_else(|| {
        Error::BadRequest(
            ErrorKind::MissingParam,
            "Missing registration token cookie.",
        )
    })?;

    let registration_token: RegistrationToken = services()
        .globals
        .validate_macaroon(token, None)
        .map_err(|_| {
        Error::BadRequest(
            ErrorKind::MissingParam,
            "Invalid registration token cookie.",
        )
    })?;

    let RegistrationInfo {
        username,
        mut displayname,
        avatar_url,
        email: _,
    } = registration_info;

    let user_id =
        UserId::parse_with_server_name(username.to_lowercase(), services().globals.server_name())
            .map_err(|_| Error::BadRequest(ErrorKind::InvalidUsername, "Invalid username."))?;

    if services().users.exists(&user_id)? {
        return Err(Error::BadRequest(
            ErrorKind::UserInUse,
            "Desired UserId is already taken.",
        ));
    }

    if services().appservice.is_exclusive_user_id(&user_id).await {
        return Err(Error::BadRequest(
            ErrorKind::Exclusive,
            "Desired UserId reserved by appservice.",
        ));
    }

    services().users.create(&user_id, None)?;
    services().users.set_password_placeholder(&user_id)?;

    if let Some(avatar_url) = avatar_url {
        let request = services().globals.default_client().get(avatar_url.as_ref());

        let res = request.send().await.map_err(|_| {
            Error::BadRequest(ErrorKind::UserInUse, "Desired UserId is already taken.")
        })?;

        let filename = avatar_url.path_segments().and_then(Iterator::last);

        let (content_type, body): (Option<headers::ContentType>, Vec<u8>) = (
            res.headers().typed_get(),
            res.bytes().await.map(Into::into).map_err(|_| {
                Error::BadRequest(ErrorKind::UserInUse, "Desired UserId is already taken.")
            })?,
        );

        let mxc = format!(
            "mxc://{}/{}",
            services().globals.server_name(),
            utils::random_string(crate::api::client_server::MXC_LENGTH)
        );

        services()
            .media
            .create(
                mxc.clone(),
                filename
                    .map(|filename| "inline; filename=".to_owned() + filename)
                    .as_deref(),
                content_type.map(|header| header.to_string()).as_deref(),
                &body,
            )
            .await?;

        services()
            .users
            .set_avatar_url(&user_id, Some(OwnedMxcUri::from(mxc)))?;
    };

    if let (Some(displayname), true) = (
        displayname.as_mut(),
        services().globals.config.enable_lightning_bolt,
    ) {
        displayname.push_str(" ⚡️");
    }

    services().users.set_displayname(&user_id, displayname)?;

    services().sso.save_claim(
        &registration_token.provider_id,
        &user_id,
        &registration_token.unique_claim,
    )?;

    services().account_data.update(
        None,
        &user_id,
        GlobalAccountDataEventType::PushRules.to_string().into(),
        &serde_json::to_value(ruma::events::push_rules::PushRulesEvent {
            content: ruma::events::push_rules::PushRulesEventContent {
                global: push::Ruleset::server_default(&user_id),
            },
        })
        .expect("PushRulesEvent should always serialize"),
    )?;

    let login_token = LoginToken::new(registration_token.provider_id, user_id);
    let redirect_uri = redirect_with_login_token(registration_token.redirect_uri, &login_token);

    Ok((
        AppendHeaders([(
            header::SET_COOKIE,
            utils::reset_cookie("sso-registration").to_string(),
        )]),
        Redirect::temporary(redirect_uri.as_str()),
    )
        .into_response())
}

fn redirect_with_login_token(mut redirect_uri: Url, login_token: &LoginToken) -> Url {
    let signed = services()
        .globals
        .sign_macaroon(login_token)
        .expect("signing macaroons should always works");

    redirect_uri
        .query_pairs_mut()
        .append_pair("loginToken", &signed);

    redirect_uri
}