pub mod attachments;
pub mod auth;
pub mod calendar;
pub mod contact_groups;
pub mod contacts;
pub mod diagnostics;
pub mod display_preferences;
pub mod drafts;
pub mod folder_counts;
pub mod folder_mgmt;
pub mod folder_preferences;
pub mod folders;
pub mod health;
pub mod identities;
pub mod mail_preferences;
pub mod messages;
pub mod notification_preferences;
pub mod pgp;
pub mod quota;
pub mod response_templates;
pub mod search;
pub mod send;
pub mod settings;
pub mod sieve;
pub mod special_folders;
pub mod system_diagnostics;
pub mod tags;
pub mod trusted_senders;
pub mod unsubscribe;

use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::sync::Arc;
use std::time::Duration;

use axum::body::Body;
use axum::extract::{ConnectInfo, DefaultBodyLimit};
use axum::http::{HeaderValue, Request, header};
use axum::response::Response;
use axum::routing::{delete, get, patch, post, put};
use axum::{Extension, Router, middleware};
use tower_governor::GovernorError;
use tower_governor::GovernorLayer;
use tower_governor::governor::GovernorConfigBuilder;
use tower_governor::key_extractor::KeyExtractor;
use tower_http::trace::TraceLayer;

use crate::auth::csrf::csrf_protection;
use crate::auth::middleware::auth_guard;
use crate::auth::session::SessionStore;
use crate::config::{AppConfig, TrustedProxy};
use crate::imap::client::ImapClient;
use crate::realtime::events::EventBus;
use crate::realtime::idle::IdleManager;
use crate::smtp::client::SmtpClient;

#[derive(Debug, Clone)]
struct ProxyAwareIpExtractor {
    trusted_proxies: Vec<TrustedProxy>,
}

impl ProxyAwareIpExtractor {
    fn is_trusted(&self, ip: &IpAddr) -> bool {
        self.trusted_proxies
            .iter()
            .any(|trusted| trusted.contains(ip))
    }

    fn forwarded_client_ip(
        &self,
        headers: &axum::http::HeaderMap,
        peer_ip: IpAddr,
    ) -> Option<IpAddr> {
        let mut current_hop = peer_ip;

        for value in headers.get_all("x-forwarded-for").iter().rev() {
            for entry in value.to_str().ok()?.split(',').rev() {
                if !self.is_trusted(&current_hop) {
                    return Some(current_hop);
                }

                current_hop = entry.trim().parse::<IpAddr>().ok()?;
            }
        }

        Some(current_hop)
    }
}

impl KeyExtractor for ProxyAwareIpExtractor {
    type Key = IpAddr;

    fn extract<T>(&self, req: &Request<T>) -> Result<Self::Key, GovernorError> {
        let peer_ip = req
            .extensions()
            .get::<ConnectInfo<SocketAddr>>()
            .map(|ci| ci.0.ip())
            .unwrap_or(IpAddr::V4(Ipv4Addr::LOCALHOST));

        if self.is_trusted(&peer_ip)
            && let Some(client_ip) = self.forwarded_client_ip(req.headers(), peer_ip)
        {
            return Ok(client_ip);
        }

        Ok(peer_ip)
    }
}

fn response_content_type_is(response: &Response, expected: &str) -> bool {
    response
        .headers()
        .get(header::CONTENT_TYPE)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.split(';').next())
        .is_some_and(|value| value.trim().eq_ignore_ascii_case(expected))
}

async fn add_security_headers(request: Request<Body>, next: middleware::Next) -> Response {
    let request_path = request.uri().path().to_owned();
    let is_api =
        request_path == "/api" || request_path.ends_with("/api") || request_path.contains("/api/");
    let mut response = next.run(request).await;
    let is_html = response_content_type_is(&response, "text/html")
        || response_content_type_is(&response, "application/xhtml+xml");
    let is_css = response_content_type_is(&response, "text/css");
    let is_pdf = response_content_type_is(&response, "application/pdf");
    let headers = response.headers_mut();

    headers.insert(
        "x-content-type-options",
        HeaderValue::from_static("nosniff"),
    );

    if is_api {
        headers.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store"));
        headers.insert(
            "cross-origin-resource-policy",
            HeaderValue::from_static("same-origin"),
        );
        if is_pdf {
            headers.insert("x-frame-options", HeaderValue::from_static("SAMEORIGIN"));
            headers.insert(
                "content-security-policy",
                HeaderValue::from_static("frame-ancestors 'self'"),
            );
        } else {
            headers.insert("x-frame-options", HeaderValue::from_static("DENY"));
            headers.insert(
                "content-security-policy",
                HeaderValue::from_static("frame-ancestors 'none'"),
            );
        }
    } else if is_html || is_css {
        headers.insert("referrer-policy", HeaderValue::from_static("no-referrer"));
    }

    if request_path.contains("/assets/") || request_path.ends_with("/favicon.svg") {
        headers.insert(
            "cross-origin-resource-policy",
            HeaderValue::from_static("same-origin"),
        );
    }

    if is_html {
        headers.insert("x-frame-options", HeaderValue::from_static("DENY"));
        headers.insert(
            "permissions-policy",
            HeaderValue::from_static("camera=(), microphone=(), geolocation=()"),
        );
        headers.insert(
            "cross-origin-opener-policy",
            HeaderValue::from_static("same-origin"),
        );
        headers.insert(
            "content-security-policy",
            HeaderValue::from_static(
                "default-src 'self'; base-uri 'self'; object-src 'none'; frame-ancestors 'none'; form-action 'self'; script-src 'self'; script-src-attr 'none'; style-src 'self'; style-src-attr 'unsafe-inline'; img-src 'self' data: blob: http: https:; font-src 'self' data:; connect-src 'self' ws: wss:; frame-src 'self' blob:",
            ),
        );
    }

    response
}

#[cfg(test)]
pub fn create_router(
    config: Arc<AppConfig>,
    store: Arc<SessionStore>,
    imap_client: Arc<dyn ImapClient>,
    smtp_client: Arc<dyn SmtpClient>,
    search_engine: Arc<crate::search::engine::SearchEngine>,
    event_bus: Arc<EventBus>,
    idle_manager: Arc<IdleManager>,
) -> Router {
    let storage =
        crate::storage::from_config(&config).expect("validated RustBox storage configuration");
    create_router_with_storage(
        config,
        RouterServices {
            store,
            imap_client,
            smtp_client,
            search_engine,
            storage,
            event_bus,
            idle_manager,
        },
    )
}

pub struct RouterServices {
    pub store: Arc<SessionStore>,
    pub imap_client: Arc<dyn ImapClient>,
    pub smtp_client: Arc<dyn SmtpClient>,
    pub search_engine: Arc<crate::search::engine::SearchEngine>,
    pub storage: Arc<dyn crate::storage::BlobStorage>,
    pub event_bus: Arc<EventBus>,
    pub idle_manager: Arc<IdleManager>,
}

pub fn create_router_with_storage(config: Arc<AppConfig>, services: RouterServices) -> Router {
    let RouterServices {
        store,
        imap_client,
        smtp_client,
        search_engine,
        storage,
        event_bus,
        idle_manager,
    } = services;
    let system_diagnostics_telemetry =
        Arc::new(system_diagnostics::SystemDiagnosticsTelemetry::new());
    let governor_conf = GovernorConfigBuilder::default()
        .key_extractor(ProxyAwareIpExtractor {
            trusted_proxies: config.parsed_trusted_proxies(),
        })
        .period(Duration::from_secs(12))
        .burst_size(5)
        .finish()
        .expect("valid governor config");

    let public_auth = Router::new()
        .route("/login", post(auth::login))
        .route("/policy", get(auth::auth_policy))
        .layer(middleware::from_fn(csrf_protection))
        .layer(GovernorLayer::new(governor_conf));

    let browser_routes = Router::new()
        .route("/accounts", get(auth::list_accounts))
        .route(
            "/accounts/{id}",
            delete(auth::remove_account).layer(middleware::from_fn(csrf_protection)),
        );

    let protected_auth = Router::new()
        .route("/session", get(auth::get_session))
        .route("/session-policy", get(auth::session_policy))
        .route("/sessions", get(auth::list_active_sessions))
        .route("/sessions/revoke-others", post(auth::revoke_other_sessions))
        .route("/sessions/{id}", delete(auth::revoke_active_session))
        .route("/logout", post(auth::logout))
        .layer(middleware::from_fn(auth_guard))
        .layer(middleware::from_fn(csrf_protection));

    let auth_router = Router::new()
        .merge(public_auth)
        .merge(browser_routes)
        .merge(protected_auth);

    let attachment_body_limit = config
        .max_attachment_size_bytes()
        .saturating_add(1024 * 1024);
    let protected_data = Router::new()
        .route(
            "/folders",
            get(folders::list_folders).post(folder_mgmt::create_folder),
        )
        .route(
            "/folders/{name}",
            patch(folder_mgmt::rename_folder).delete(folder_mgmt::delete_folder),
        )
        .route(
            "/folders/{name}/subscribe",
            patch(folder_mgmt::subscribe_folder),
        )
        .route(
            "/folders/{name}/mark-all-read",
            post(folder_mgmt::mark_all_read),
        )
        .route("/folders/{name}/empty", post(folder_mgmt::empty_folder))
        .route(
            "/folders/{folder}/import-eml",
            post(messages::import_eml).layer(DefaultBodyLimit::max(52 * 1024 * 1024)),
        )
        .route("/folders/{folder}/messages", get(messages::list_messages))
        .route(
            "/messages/{folder}/{uid}",
            get(messages::get_message).delete(messages::delete_message_handler),
        )
        .route(
            "/messages/{folder}/{uid}/flags",
            patch(messages::update_flags),
        )
        .route(
            "/messages/{folder}/{uid}/attachments/{attachment_id}",
            get(messages::download_attachment),
        )
        .route(
            "/messages/{folder}/{uid}/source",
            get(messages::download_eml_source),
        )
        .route("/messages/move", post(messages::move_message_handler))
        .route("/messages/copy", post(messages::copy_message_handler))
        .route("/messages/send", post(send::send_message_handler))
        .route(
            "/messages/{folder}/{uid}/unsubscribe",
            post(unsubscribe::unsubscribe_message),
        )
        .route(
            "/drafts",
            get(drafts::list_drafts_handler).post(drafts::upsert_draft_handler),
        )
        .route(
            "/drafts/{id}",
            get(drafts::get_draft_handler).delete(drafts::delete_draft_handler),
        )
        .route(
            "/drafts/{draft_id}/attachments",
            post(attachments::upload_attachment)
                .layer(DefaultBodyLimit::max(attachment_body_limit)),
        )
        .route(
            "/drafts/{draft_id}/attachments/{attachment_id}",
            delete(attachments::delete_attachment),
        )
        .route(
            "/drafts/{draft_id}/attachments/{attachment_id}/content",
            get(attachments::get_attachment_content),
        )
        .route("/search", get(search::search_messages))
        .route(
            "/response-templates",
            get(response_templates::list).post(response_templates::create),
        )
        .route(
            "/response-templates/{id}",
            put(response_templates::update).delete(response_templates::delete),
        )
        .route(
            "/search/saved",
            get(search::list_saved_searches).post(search::create_saved_search),
        )
        .route(
            "/search/saved/{id}",
            axum::routing::put(search::update_saved_search).delete(search::delete_saved_search),
        )
        .route(
            "/contact-groups",
            get(contact_groups::list_groups_handler).post(contact_groups::create_group_handler),
        )
        .route(
            "/contact-groups/{id}",
            put(contact_groups::update_group_handler).delete(contact_groups::delete_group_handler),
        )
        .route(
            "/contact-groups/{id}/members",
            get(contact_groups::list_members_handler).post(contact_groups::add_member_handler),
        )
        .route(
            "/contact-groups/{id}/members/{contact_id}",
            delete(contact_groups::remove_member_handler),
        )
        .route(
            "/tags",
            get(tags::list_tags_handler).post(tags::create_tag_handler),
        )
        .route("/tags/reconcile", post(tags::reconcile_tags_handler))
        .route(
            "/tags/{id}",
            put(tags::update_tag_handler).delete(tags::delete_tag_handler),
        )
        .route(
            "/tags/{id}/messages",
            post(tags::tag_message_handler).get(tags::list_tag_messages_handler),
        )
        .route("/tags/{id}/messages/bulk", post(tags::bulk_tag_handler))
        .route(
            "/tags/{id}/messages/{folder}/{uid}",
            delete(tags::untag_message_handler),
        )
        .route(
            "/messages/{folder}/{uid}/tags",
            get(tags::get_message_tags_handler),
        )
        .route(
            "/contacts",
            get(contacts::list_contacts_handler).post(contacts::create_contact_handler),
        )
        .route(
            "/contacts/{id}/photo",
            get(contacts::get_contact_photo_handler)
                .post(contacts::upload_contact_photo_handler)
                .delete(contacts::delete_contact_photo_handler)
                .layer(DefaultBodyLimit::max(3 * 1024 * 1024)),
        )
        .route(
            "/contacts/duplicates",
            get(contacts::duplicate_contacts_handler),
        )
        .route("/contacts/merge", post(contacts::merge_contacts_handler))
        .route("/contacts/export", get(contacts::export_contacts_handler))
        .route(
            "/contacts/import",
            post(contacts::import_contacts_handler).layer(DefaultBodyLimit::max(11 * 1024 * 1024)),
        )
        .route(
            "/contacts/autocomplete/all",
            get(contacts::autocomplete_all_handler),
        )
        .route(
            "/contacts/autocomplete",
            get(contacts::autocomplete_handler),
        )
        .route(
            "/contacts/{id}",
            get(contacts::get_contact_handler).delete(contacts::delete_contact_handler),
        )
        .route(
            "/contacts/{contact_id}/groups",
            get(contact_groups::list_contact_groups_handler),
        )
        .route(
            "/contacts/{id}/export",
            get(contacts::export_single_contact_handler),
        )
        .route(
            "/identities",
            get(identities::list_identities_handler)
                .post(identities::create_identity_handler)
                .layer(DefaultBodyLimit::max(4 * 1024 * 1024)),
        )
        .route(
            "/identities/{id}",
            get(identities::get_identity_handler)
                .put(identities::update_identity_handler)
                .delete(identities::delete_identity_handler)
                .layer(DefaultBodyLimit::max(4 * 1024 * 1024)),
        )
        .route(
            "/pgp/keys",
            get(pgp::list_keys)
                .post(pgp::import_key)
                .layer(DefaultBodyLimit::max(3 * 1024 * 1024)),
        )
        .route("/pgp/keys/generate", post(pgp::generate_key))
        .route("/pgp/keys/{fingerprint}/export", get(pgp::export_key))
        .route(
            "/pgp/keys/{fingerprint}/test-passphrase",
            post(pgp::test_passphrase),
        )
        .route(
            "/pgp/keys/{fingerprint}/change-passphrase",
            post(pgp::change_passphrase),
        )
        .route("/pgp/keys/{fingerprint}/trust", patch(pgp::set_trust))
        .route("/pgp/keys/{fingerprint}", delete(pgp::delete_key))
        .route(
            "/pgp/preferences",
            get(pgp::get_preferences).put(pgp::update_preferences),
        )
        .route("/pgp/discover", post(pgp::discover_key))
        .route("/pgp/recipients", post(pgp::recipient_status))
        .route("/pgp/verify", post(pgp::verify))
        .route(
            "/pgp/decrypt",
            post(pgp::decrypt).layer(DefaultBodyLimit::max(12 * 1024 * 1024)),
        )
        .route(
            "/pgp/messages/{folder}/{uid}/process",
            post(pgp::process_message),
        )
        .route(
            "/settings/display",
            get(display_preferences::get_display_preferences)
                .put(display_preferences::update_display_preferences),
        )
        .route(
            "/settings/special-folders",
            get(special_folders::get_special_folders).put(special_folders::update_special_folders),
        )
        .route(
            "/settings/notifications",
            get(notification_preferences::get_notification_preferences)
                .put(notification_preferences::update_notification_preferences),
        )
        .route(
            "/settings/mailbox",
            get(mail_preferences::get_mail_preferences)
                .put(mail_preferences::update_mail_preferences),
        )
        .route(
            "/settings/trusted-senders",
            get(trusted_senders::list_trusted_senders).post(trusted_senders::add_trusted_sender),
        )
        .route(
            "/settings/trusted-senders/{email}",
            delete(trusted_senders::remove_trusted_sender),
        )
        .route(
            "/folders/{folder}/preferences",
            get(folder_preferences::get_folder_preferences)
                .put(folder_preferences::update_folder_preferences),
        )
        .route(
            "/calendar/events",
            get(calendar::list_events).post(calendar::create_event),
        )
        .route(
            "/calendar/events/import-ics",
            post(calendar::import_ics).layer(DefaultBodyLimit::max(11 * 1024 * 1024)),
        )
        .route("/calendar/events/export-ics", get(calendar::export_ics))
        .route(
            "/calendar/events/{id}",
            get(calendar::get_event)
                .put(calendar::update_event)
                .delete(calendar::delete_event),
        )
        .route(
            "/calendar/settings",
            get(calendar::get_calendar_settings).put(calendar::update_calendar_settings),
        )
        .route(
            "/calendar/meeting-templates",
            get(calendar::list_meeting_templates).post(calendar::create_meeting_template),
        )
        .route(
            "/calendar/meeting-templates/{id}",
            delete(calendar::delete_meeting_template),
        )
        .route("/quota", get(quota::get_quota))
        .route("/settings/reset", post(settings::reset_all_settings));

    let diagnostics_routes = Router::new()
        .route("/diagnostics/imap", get(diagnostics::imap_diagnostics))
        .route("/diagnostics/smtp", get(diagnostics::smtp_diagnostics))
        .route("/diagnostics/sync", get(diagnostics::sync_diagnostics))
        .route(
            "/diagnostics/imap/reconnect",
            post(diagnostics::reconnect_imap),
        )
        .route("/diagnostics/search", get(diagnostics::search_diagnostics))
        .route(
            "/diagnostics/search/rebuild",
            post(diagnostics::rebuild_search_index),
        );

    let sieve_routes = Router::new()
        .route("/sieve/status", get(sieve::status))
        .route("/sieve/scripts", get(sieve::list_scripts))
        .route(
            "/sieve/scripts/{name}",
            get(sieve::get_script)
                .put(sieve::put_script)
                .delete(sieve::delete_script),
        )
        .route(
            "/sieve/scripts/{name}/activate",
            post(sieve::activate_script),
        )
        .route("/sieve/scripts/{name}/rename", post(sieve::rename_script))
        .route("/sieve/deactivate", post(sieve::deactivate_script))
        .route("/sieve/check", post(sieve::check_script));

    let protected_data = if config.diagnostics_enabled {
        protected_data.merge(diagnostics_routes)
    } else {
        protected_data
    };
    let protected_data = if config.sieve_enabled {
        protected_data.merge(sieve_routes)
    } else {
        protected_data
    }
    .layer(middleware::from_fn(auth_guard))
    .layer(middleware::from_fn(csrf_protection));

    let ws_route = Router::new().route("/ws", get(crate::realtime::ws::ws_handler));

    let mut api_router = Router::new()
        .route("/health", get(health::health_check))
        .route("/ready", get(health::readiness_check))
        .nest("/auth", auth_router)
        .merge(ws_route)
        .merge(protected_data);

    if config.diagnostics_endpoint.enabled {
        let diagnostics_path = config
            .diagnostics_endpoint
            .path
            .trim()
            .strip_prefix("/api")
            .unwrap_or(config.diagnostics_endpoint.path.trim());
        api_router = api_router.route(
            diagnostics_path,
            get(system_diagnostics::system_diagnostics),
        );
    }

    let api_router = api_router.fallback(|| async { axum::http::StatusCode::NOT_FOUND });

    let inner = Router::new()
        .nest("/api", api_router)
        .route("/", get(crate::ui::serve_ui))
        .fallback(crate::ui::serve_ui);

    let router = match config.normalized_base_path() {
        "/" => inner,
        base_path => {
            let base_path_slash = format!("{}/", base_path.trim_end_matches('/'));
            Router::new()
                .route(&base_path_slash, get(crate::ui::serve_ui))
                .nest(base_path, inner)
        }
    };

    let router = if config.diagnostics_endpoint.enabled {
        let system_diagnostics_dependencies =
            Arc::new(system_diagnostics::SystemDiagnosticsDependencies::new(
                config.clone(),
                store.clone(),
                search_engine.clone(),
                storage.clone(),
                idle_manager.clone(),
                system_diagnostics_telemetry.clone(),
            ));
        router
            .layer(middleware::from_fn(system_diagnostics::observe_requests))
            .layer(Extension(system_diagnostics_dependencies))
            .layer(Extension(system_diagnostics_telemetry))
    } else {
        router
    };

    let router = router
        .layer(Extension(idle_manager))
        .layer(Extension(event_bus))
        .layer(Extension(smtp_client))
        .layer(Extension(search_engine))
        .layer(Extension(storage))
        .layer(Extension(imap_client))
        .layer(Extension(store))
        .layer(Extension(config.clone()))
        .layer(TraceLayer::new_for_http());

    if config.security_headers_enabled {
        router.layer(middleware::from_fn(add_security_headers))
    } else {
        router
    }
}

#[cfg(test)]
mod tests;