// Copyright 2024, 2025 New Vector Ltd.
//
// SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
// Please see LICENSE files in the repository root for full details.

use std::{net::IpAddr, sync::Arc};

use axum::{
    extract::{Form, State},
    response::{Html, IntoResponse, Response},
};
use axum_extra::extract::Query;
use mas_axum_utils::{
    InternalError, SessionInfoExt,
    cookies::CookieJar,
    csrf::{CsrfExt as _, ProtectedForm},
};
use mas_data_model::{BoxClock, BoxRng, SiteConfig};
use mas_matrix::HomeserverConnection;
use mas_policy::Policy;
use mas_router::{PasswordRegister, UpstreamOAuth2Authorize, UrlBuilder};
use mas_storage::{BoxRepository, RepositoryAccess};
use mas_templates::{
    FieldError, FormError, FormState, RegisterContext, RegisterFormField, TemplateContext,
    Templates, ToFormState,
};
use serde::{Deserialize, Serialize};

use super::shared::OptionalPostAuthAction;
use crate::{BoundActivityTracker, PreferredLanguage};

mod cookie;
pub(crate) mod password;
pub(crate) mod steps;

pub use self::cookie::UserRegistrationSessions as UserRegistrationSessionsCookie;

#[derive(Debug, Deserialize, Serialize)]
pub(crate) struct UsernameForm {
    username: String,
    #[serde(flatten)]
    action: OptionalPostAuthAction,
}

impl ToFormState for UsernameForm {
    type Field = RegisterFormField;
}

#[allow(clippy::too_many_arguments)]
pub(super) async fn validate_registration_username(
    state: &mut FormState<RegisterFormField>,
    username: &str,
    user_agent: Option<String>,
    ip_address: Option<IpAddr>,
    repo: &mut BoxRepository,
    policy: &mut Policy,
    homeserver: &Arc<dyn HomeserverConnection>,
    http_client: &reqwest::Client,
    site_config: &SiteConfig,
) -> Result<(), InternalError> {
    let mut homeserver_denied_username = false;
    if username.is_empty() {
        state.add_error_on_field(RegisterFormField::Username, FieldError::Required);
    } else if repo.user().exists(username).await? {
        state.add_error_on_field(RegisterFormField::Username, FieldError::Exists);
    } else if !homeserver
        .is_localpart_available(username)
        .await
        .map_err(InternalError::from_anyhow)?
    {
        tracing::warn!(username, "Homeserver denied username provided by user");
        homeserver_denied_username = true;
    }

    let res = policy
        .evaluate_register(mas_policy::RegisterInput {
            registration_method: mas_policy::RegistrationMethod::Password,
            username,
            email: None,
            requester: mas_policy::Requester {
                ip_address,
                user_agent,
            },
        })
        .await?;

    for violation in res.violations {
        match violation.field.as_deref() {
            Some("username") => {
                homeserver_denied_username = false;
                state.add_error_on_field(
                    RegisterFormField::Username,
                    FieldError::Policy {
                        code: violation.variant.map(|c| c.as_str()),
                        message: violation.msg,
                    },
                );
            }
            _ => state.add_error_on_form(FormError::Policy {
                code: violation.variant.map(|c| c.as_str()),
                message: violation.msg,
            }),
        }
    }

    if homeserver_denied_username {
        state.add_error_on_field(RegisterFormField::Username, FieldError::Exists);
    }

    if !username.is_empty() && state.is_valid() {
        match crate::postnumber::check(
            http_client,
            site_config.postnumber_validation.as_ref(),
            username,
        )
        .await
        {
            Ok(Some(crate::postnumber::PostnumberOutcome::Allowed) | None) => {}
            Ok(Some(crate::postnumber::PostnumberOutcome::Reserved)) => {
                state.add_error_on_field(
                    RegisterFormField::Username,
                    FieldError::Policy {
                        code: Some("postnumber-reserved"),
                        message: "This username is reserved".to_owned(),
                    },
                );
            }
            Err(e) => {
                tracing::error!(
                    error = &e as &dyn std::error::Error,
                    "postnumber resolver error"
                );
                state.add_error_on_form(FormError::Internal);
            }
        }
    }

    Ok(())
}

async fn render(
    locale: mas_i18n::DataLocale,
    mut ctx: RegisterContext,
    action: OptionalPostAuthAction,
    csrf_token: impl ToString,
    repo: &mut BoxRepository,
    templates: &Templates,
) -> Result<String, InternalError> {
    let post_action = action
        .load_context(repo)
        .await
        .map_err(InternalError::from_anyhow)?;
    if let Some(action) = post_action {
        ctx = ctx.with_post_action(action);
    }

    Ok(templates.render_register(&ctx.with_csrf(csrf_token).with_language(locale))?)
}

#[tracing::instrument(name = "handlers.views.register.get", skip_all)]
pub(crate) async fn get(
    mut rng: BoxRng,
    clock: BoxClock,
    PreferredLanguage(locale): PreferredLanguage,
    State(templates): State<Templates>,
    State(url_builder): State<UrlBuilder>,
    State(site_config): State<SiteConfig>,
    mut repo: BoxRepository,
    activity_tracker: BoundActivityTracker,
    Query(query): Query<OptionalPostAuthAction>,
    cookie_jar: CookieJar,
) -> Result<Response, InternalError> {
    let (csrf_token, cookie_jar) = cookie_jar.csrf_token(&clock, &mut rng);
    let (session_info, cookie_jar) = cookie_jar.session_info();

    let maybe_session = session_info.load_active_session(&mut repo).await?;

    if let Some(session) = maybe_session {
        activity_tracker
            .record_browser_session(&clock, &session)
            .await;

        let reply = query.go_next(&url_builder);
        return Ok((cookie_jar, reply).into_response());
    }

    let providers = repo.upstream_oauth_provider().all_enabled().await?;

    // If password-based login is disabled, and there is only one upstream provider,
    // we can directly start an authorization flow
    if !site_config.password_registration_enabled && providers.len() == 1 {
        let provider = providers.into_iter().next().unwrap();

        let mut destination = UpstreamOAuth2Authorize::new(provider.id);

        if let Some(action) = query.post_auth_action {
            destination = destination.and_then(action);
        }

        return Ok((cookie_jar, url_builder.redirect(&destination)).into_response());
    }

    // If password-based registration is enabled and there are no upstream
    // providers, we redirect to the password registration page
    if site_config.password_registration_enabled && providers.is_empty() {
        let mut destination = PasswordRegister::default();

        if let Some(action) = query.post_auth_action {
            destination = destination.and_then(action);
        }

        return Ok((cookie_jar, url_builder.redirect(&destination)).into_response());
    }

    let content = render(
        locale,
        RegisterContext::new(providers),
        query,
        csrf_token.form_value(),
        &mut repo,
        &templates,
    )
    .await?;

    Ok((cookie_jar, Html(content)).into_response())
}

#[tracing::instrument(name = "handlers.views.register.post", skip_all)]
pub(crate) async fn post(
    mut rng: BoxRng,
    clock: BoxClock,
    PreferredLanguage(locale): PreferredLanguage,
    State(templates): State<Templates>,
    State(url_builder): State<UrlBuilder>,
    State(site_config): State<SiteConfig>,
    State(homeserver): State<Arc<dyn HomeserverConnection>>,
    State(http_client): State<reqwest::Client>,
    mut policy: Policy,
    mut repo: BoxRepository,
    activity_tracker: BoundActivityTracker,
    user_agent: Option<axum_extra::typed_header::TypedHeader<headers::UserAgent>>,
    cookie_jar: CookieJar,
    Form(form): Form<ProtectedForm<UsernameForm>>,
) -> Result<Response, InternalError> {
    let user_agent = user_agent.map(|ua| ua.as_str().to_owned());

    let form = cookie_jar.verify_form(&clock, form)?;
    let (csrf_token, cookie_jar) = cookie_jar.csrf_token(&clock, &mut rng);

    if !site_config.password_registration_enabled {
        return Ok(url_builder
            .redirect(&mas_router::Login::from(form.action.post_auth_action))
            .into_response());
    }

    let providers = repo.upstream_oauth_provider().all_enabled().await?;
    let mut state = form.to_form_state();
    validate_registration_username(
        &mut state,
        &form.username,
        user_agent,
        activity_tracker.ip(),
        &mut repo,
        &mut policy,
        &homeserver,
        &http_client,
        &site_config,
    )
    .await?;

    if !state.is_valid() {
        let content = render(
            locale,
            RegisterContext::new(providers).with_form_state(state),
            form.action,
            csrf_token.form_value(),
            &mut repo,
            &templates,
        )
        .await?;

        return Ok((cookie_jar, Html(content)).into_response());
    }

    let mut destination = PasswordRegister::default().with_username(form.username);
    if let Some(action) = form.action.post_auth_action {
        destination = destination.and_then(action);
    }

    Ok((cookie_jar, url_builder.redirect(&destination)).into_response())
}
