use std::{future::Future, net::SocketAddr, str::FromStr};

use airgap::AirgapDetection;
use axum::{
    Router,
    extract::{DefaultBodyLimit, FromRef},
    http::StatusCode,
    middleware,
    response::IntoResponse,
    routing::any,
    serve::ListenerExt,
};
use hyper::http::{HeaderName, HeaderValue, header};
use sqlx::{
    Sqlite, SqliteConnection, SqlitePool,
    sqlite::{SqliteConnectOptions, SqliteJournalMode},
};
use tokio::{net::TcpListener, signal};
use tower_http::{
    set_header::SetResponseHeaderLayer,
    trace::{self, TraceLayer},
};
use tracing::{Level, info, trace, warn};
use utoipa::{
    Modify, OpenApi,
    openapi::{
        path::Operation,
        security::{ApiKey, ApiKeyValue, SecurityScheme},
    },
};
use utoipa_axum::router::OpenApiRouter;
#[cfg(feature = "openapi")]
use utoipa_swagger_ui::SwaggerUi;

pub mod airgap;
pub mod app_error;
pub mod audit_log;
pub mod authentication;
pub mod committee_session;
pub mod data_entry;
pub mod document;
pub mod election;
pub mod eml;
mod error;
pub mod files;
#[cfg(feature = "dev-database")]
pub mod fixtures;
pub mod investigation;
pub mod pdf_gen;
pub mod polling_station;
pub mod report;
pub mod summary;
#[cfg(feature = "dev-database")]
pub mod test_data_gen;
pub mod zip;

pub use app_error::AppError;
pub use error::{APIError, ErrorResponse};

use crate::app_error::{DatabaseErrorWithPath, DatabaseMigrationErrorWithPath};
use crate::authentication::SECURITY_SCHEME_NAME;

/// Maximum size of the request body in megabytes.
pub const MAX_BODY_SIZE_MB: usize = 12;

/// Extension trait for SqlitePool to add begin_immediate functionality
pub trait SqlitePoolExt {
    /// Acquire a connection and start a transaction with "BEGIN IMMEDIATE",
    /// which is needed to prevent database is busy errors.
    /// See https://sqlite.org/isolation.html for details.
    fn begin_immediate(
        &self,
    ) -> impl Future<Output = Result<sqlx::Transaction<'_, Sqlite>, sqlx::Error>> + Send;
}

impl SqlitePoolExt for SqlitePool {
    async fn begin_immediate(&self) -> Result<sqlx::Transaction<'_, Sqlite>, sqlx::Error> {
        self.begin_with("BEGIN IMMEDIATE").await
    }
}

#[derive(FromRef, Clone)]
pub struct AppState {
    pool: SqlitePool,
    airgap_detection: AirgapDetection,
}

pub fn get_scopes_from_operation(operation: &Operation) -> Option<Vec<String>> {
    let security_reqs = operation.security.as_ref()?;

    let scopes: Vec<String> = security_reqs
        .iter()
        .filter_map(|req| {
            // Serialization to access private BTreeMap values of SecurityRequirement
            // Proposed change upstream: https://github.com/juhaku/utoipa/pull/1494
            serde_json::to_value(req).ok().and_then(|v| {
                v.as_object()?
                    .get(SECURITY_SCHEME_NAME)?
                    .as_array()
                    .cloned()
            })
        })
        .flat_map(|arr| arr.into_iter().filter_map(|v| v.as_str().map(String::from)))
        .collect();

    Some(scopes)
}

pub fn openapi_router() -> OpenApiRouter<AppState> {
    #[derive(utoipa::OpenApi)]
    struct ApiDoc;

    struct SecurityAddon;

    impl utoipa::Modify for SecurityAddon {
        fn modify(&self, openapi: &mut utoipa::openapi::OpenApi) {
            if let Some(components) = openapi.components.as_mut() {
                components.add_security_scheme(
                    "cookie_auth",
                    SecurityScheme::ApiKey(ApiKey::Cookie(ApiKeyValue::new(String::from(
                        crate::authentication::SESSION_COOKIE_NAME,
                    )))),
                )
            }

            // Add scopes to operation summaries
            for path_item in openapi.paths.paths.values_mut() {
                for operation in [
                    &mut path_item.get,
                    &mut path_item.post,
                    &mut path_item.put,
                    &mut path_item.delete,
                    &mut path_item.patch,
                ]
                .into_iter()
                .flatten()
                {
                    let scopes = get_scopes_from_operation(operation);

                    if let Some(scopes) = scopes {
                        let extra = scopes.join(", ");
                        operation.summary = Some(match &operation.summary {
                            Some(existing) => format!("{} ({})", existing, extra),
                            None => extra,
                        });
                    }
                }
            }
        }
    }

    let mut router = build_routes(ApiDoc::openapi());

    // Modify schema to add security schemes after building routes
    SecurityAddon.modify(router.get_openapi_mut());

    router
}

fn build_routes(doc: utoipa::openapi::OpenApi) -> OpenApiRouter<AppState> {
    let router = OpenApiRouter::with_openapi(doc)
        .merge(audit_log::router())
        .merge(authentication::router())
        .merge(authentication::user_router())
        .merge(committee_session::router())
        .merge(data_entry::router())
        .merge(election::router())
        .merge(polling_station::router())
        .merge(report::router())
        .merge(document::router())
        .merge(investigation::router());

    #[cfg(feature = "dev-database")]
    let router = router.merge(test_data_gen::router());

    router
}

fn axum_router_from_openapi(router: OpenApiRouter<AppState>) -> Router<AppState> {
    // Serve the OpenAPI documentation at /api-docs if the openapi feature is enabled
    #[cfg(feature = "openapi")]
    let router = {
        let (router, openapi) = router.split_for_parts();
        router.merge(SwaggerUi::new("/api-docs").url("/api-docs/openapi.json", openapi))
    };

    // Serve without OpenAPI documentation if the openapi feature is disabled
    #[cfg(not(feature = "openapi"))]
    let router = Router::from(router);

    // Add fallback for unmatched /api/* routes to return 404
    // This must come before the frontend memory-serve routes
    async fn api_fallback_handler() -> impl IntoResponse {
        (StatusCode::NOT_FOUND, "Not Found")
    }
    router
        .route("/api", any(api_fallback_handler))
        .route("/api/", any(api_fallback_handler))
        .route("/api/{*path}", any(api_fallback_handler))
}

/// Add middleware to trace all HTTP requests and extend the user's session lifetime if needed
fn add_middleware(router: Router<AppState>, state: &AppState) -> Router<AppState> {
    router
        .layer(middleware::map_response(error::map_error_response))
        .layer(DefaultBodyLimit::max(1024 * 1024 * MAX_BODY_SIZE_MB))
        .layer(
            TraceLayer::new_for_http()
                .make_span_with(trace::DefaultMakeSpan::new().level(Level::INFO))
                .on_response(trace::DefaultOnResponse::new().level(Level::INFO)),
        )
        .layer(middleware::map_response_with_state(
            state.clone(),
            audit_log::log_error,
        ))
        .layer(middleware::map_response_with_state(
            state.clone(),
            authentication::extend_session,
        ))
        // Caution: make sure "inject_user" is added after "extend_session"
        .layer(middleware::map_request_with_state(
            state.clone(),
            authentication::inject_user,
        ))
        .layer(middleware::from_fn_with_state(
            state.clone(),
            airgap::block_request_on_airgap_violation,
        ))
        // Set Cache-Control header to prevent caching of API responses
        .layer(SetResponseHeaderLayer::overriding(
            header::CACHE_CONTROL,
            HeaderValue::from_static("no-store"),
        ))
}

/// Add the memory-serve router to serve the frontend (if the memory-serve feature is enabled)
/// Note that memory-serve includes its own Cache-Control header.
#[cfg(feature = "memory-serve")]
fn add_frontend_memory_serve(router: Router<AppState>) -> Router<AppState> {
    router.merge(
        memory_serve::from_local_build!("frontend")
            .index_file(Some("/index.html"))
            .fallback(Some("/index.html"))
            .fallback_status(StatusCode::OK)
            .into_router()
            // Add Referrer-Policy, Permissions-Policy and Content-Security-Policy headers.
            // These are only needed on HTML documents, not API requests.
            // From https://cheatsheetseries.owasp.org/cheatsheets/REST_Security_Cheat_Sheet.html#security-headers:
            // "The headers below are only intended to provide additional security when responses are rendered as HTML."
            .layer(SetResponseHeaderLayer::overriding(
                header::REFERRER_POLICY,
                HeaderValue::from_static("no-referrer"),
            ))
            .layer(SetResponseHeaderLayer::overriding(
                HeaderName::from_static("permissions-policy"),
                HeaderValue::from_static(
                    "accelerometer=(), autoplay=(), camera=(), cross-origin-isolated=(), \
                     display-capture=(), encrypted-media=(), fullscreen=(), geolocation=(), \
                     gyroscope=(), keyboard-map=(), magnetometer=(), microphone=(), midi=(), \
                     payment=(), picture-in-picture=(), publickey-credentials-get=(), \
                     screen-wake-lock=(), sync-xhr=(self), usb=(), web-share=(), \
                     xr-spatial-tracking=(), clipboard-read=(), clipboard-write=(), gamepad=(), \
                     hid=(), idle-detection=(), interest-cohort=(), serial=(), unload=()",
                ),
            ))
            .layer(SetResponseHeaderLayer::overriding(
                header::CONTENT_SECURITY_POLICY,
                HeaderValue::from_static("default-src 'self'; img-src 'self' data:"),
            )),
    )
}

/// Add memory-serve router to serve Storybook
#[cfg(feature = "storybook")]
fn add_storybook_memory_serve(router: Router<AppState>) -> Router<AppState> {
    router
        .nest(
            "/storybook/",
            memory_serve::from_local_build!("storybook")
                .index_file(Some("/index.html"))
                .into_router()
                .layer(SetResponseHeaderLayer::overriding(
                    header::X_FRAME_OPTIONS,
                    HeaderValue::from_static("sameorigin"),
                )),
        )
        // Workaround for https://github.com/storybookjs/storybook/issues/32428
        .route(
            "/vite-inject-mocker-entry.js",
            axum::routing::get(|| async {
                axum::response::Redirect::temporary("/storybook/vite-inject-mocker-entry.js")
            }),
        )
}

/// Add headers for security hardening
/// Best practices according to the OWASP Secure Headers Project, https://owasp.org/www-project-secure-headers/
fn add_security_headers(router: Router<AppState>) -> Router<AppState> {
    let security_headers_service = tower::ServiceBuilder::new()
        .layer(SetResponseHeaderLayer::if_not_present(
            header::X_FRAME_OPTIONS,
            HeaderValue::from_static("deny"),
        ))
        .layer(SetResponseHeaderLayer::overriding(
            header::X_CONTENT_TYPE_OPTIONS,
            HeaderValue::from_static("nosniff"),
        ))
        .layer(SetResponseHeaderLayer::overriding(
            HeaderName::from_static("x-permitted-cross-domain-policies"),
            HeaderValue::from_static("none"),
        ))
        .layer(SetResponseHeaderLayer::overriding(
            HeaderName::from_static("cross-origin-resource-policy"),
            HeaderValue::from_static("same-origin"),
        ))
        .layer(SetResponseHeaderLayer::overriding(
            HeaderName::from_static("cross-origin-embedder-policy"),
            HeaderValue::from_static("require-corp"),
        ))
        .layer(SetResponseHeaderLayer::overriding(
            HeaderName::from_static("cross-origin-opener-policy"),
            HeaderValue::from_static("same-origin"),
        ));
    router.layer(security_headers_service)
}

/// Complete Axum router for the application
pub fn router(pool: SqlitePool, airgap_detection: AirgapDetection) -> Result<Router, AppError> {
    let router = axum_router_from_openapi(openapi_router());
    let state = AppState {
        pool,
        airgap_detection,
    };
    let router = add_middleware(router, &state);
    #[cfg(feature = "memory-serve")]
    let router = add_frontend_memory_serve(router);
    #[cfg(feature = "storybook")]
    let router = add_storybook_memory_serve(router);
    let router = add_security_headers(router);
    let router = router.with_state(state);
    Ok(router)
}

/// Start the API server on the given port, using the given database pool.
pub async fn start_server(
    pool: SqlitePool,
    listener: TcpListener,
    enable_airgap_detection: bool,
) -> Result<(), AppError> {
    info!(
        "Starting Abacus GR26 (version {})",
        env!("ABACUS_GIT_VERSION")
    );
    let airgap_detection = if enable_airgap_detection {
        info!("Airgap detection is enabled, starting airgap detection task...");

        AirgapDetection::start(pool.clone()).await
    } else {
        warn!("Airgap detection is disabled, this is not allowed in production.");

        AirgapDetection::nop()
    };

    let app = router(pool.clone(), airgap_detection)?;

    info!("Starting Abacus on http://{}", listener.local_addr()?);
    let listener = listener.tap_io(|tcp_stream| {
        if let Err(err) = tcp_stream.set_nodelay(true) {
            trace!("failed to set TCP_NODELAY on incoming connection: {err:#}");
        }
    });

    axum::serve(
        listener,
        app.into_make_service_with_connect_info::<SocketAddr>(),
    )
    .with_graceful_shutdown(shutdown_signal())
    .await?;

    // close the database, this will flush the shm and wal files for sqlite
    pool.close().await;

    info!("Abacus has shut down gracefully.");

    Ok(())
}

/// Graceful shutdown, useful for Docker containers.
///
/// Copied from the
/// [axum graceful-shutdown example](https://github.com/tokio-rs/axum/blob/6318b57fda6b524b4d3c7909e07946e2b246ebd2/examples/graceful-shutdown/src/main.rs)
/// (under the MIT license).
pub async fn shutdown_signal() {
    let ctrl_c = async {
        signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("failed to install terminate signal handler")
            .recv()
            .await;
    };

    #[cfg(unix)]
    let quit = async {
        signal::unix::signal(signal::unix::SignalKind::quit())
            .expect("failed to install quit signal handler")
            .recv()
            .await;
    };

    #[cfg(windows)]
    let terminate = async {
        signal::windows::ctrl_close()
            .expect("failed to install CTRL_CLOSE handler")
            .recv()
            .await
    };

    #[cfg(windows)]
    let quit = async {
        signal::windows::ctrl_shutdown()
            .expect("failed to install CTRL_SHUTDOWN handler")
            .recv()
            .await
    };

    #[cfg(not(any(unix, windows)))]
    let terminate = std::future::pending::<()>();

    #[cfg(not(any(unix, windows)))]
    let quit = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
        _ = quit => {},
    }
}

/// Log that the application started and thereby check that the database is writeable
/// This maps common sqlite errors to an AppError
async fn log_app_started(conn: &mut SqliteConnection, db_path: &str) -> Result<(), AppError> {
    Ok(audit_log::create(
        conn,
        &audit_log::AuditEvent::ApplicationStarted(audit_log::ApplicationStartedDetails {
            version: env!("ABACUS_GIT_VERSION").to_string(),
            commit: env!("ABACUS_GIT_REV").to_string(),
        }),
        None,
        None,
        None,
    )
    .await
    .map_err(DatabaseErrorWithPath::with_path(db_path))?)
}

/// Create a SQLite database if needed, then connect to it and run migrations.
/// Return a connection pool.
pub async fn create_sqlite_pool(
    database: &str,
    #[cfg(feature = "dev-database")] reset_database: bool,
    #[cfg(feature = "dev-database")] seed_data: bool,
) -> Result<SqlitePool, AppError> {
    let db = format!("sqlite://{database}");
    let opts = SqliteConnectOptions::from_str(&db)?
        .create_if_missing(true)
        .journal_mode(SqliteJournalMode::Wal)
        .busy_timeout(std::time::Duration::from_secs(20));

    #[cfg(feature = "dev-database")]
    if reset_database {
        // remove the file, ignoring any errors that occurred (such as the file not existing)
        let _ = tokio::fs::remove_file(opts.get_filename()).await;
        info!("removed database file");
    }

    let pool = SqlitePool::connect_with(opts).await?;

    // run database migrations, this creates the necessary tables if they don't exist yet
    sqlx::migrate!()
        .run(&pool)
        .await
        .map_err(DatabaseMigrationErrorWithPath::with_path(database))?;

    #[cfg(feature = "dev-database")]
    if seed_data {
        fixtures::seed_fixture_data(&pool).await?;
    }

    // log startup event and verify the database is writeable
    let mut connection = pool.acquire().await?;
    log_app_started(&mut connection, database).await?;

    Ok(pool)
}

#[cfg(test)]
mod test {
    use sqlx::SqlitePool;
    use test_log::test;
    use tokio::net::TcpListener;

    use crate::{AppError, create_sqlite_pool};

    use super::start_server;

    async fn run_server_test<F, Fut>(pool: SqlitePool, test_fn: F)
    where
        F: FnOnce(String) -> Fut,
        Fut: Future<Output = ()>,
    {
        // Setup: create server on random port
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let base_url = format!("http://{addr}");

        // Start server in background task
        let server_task = tokio::spawn(async move {
            start_server(pool, listener, false).await.unwrap();
        });

        // Run the test
        test_fn(base_url).await;

        // Cleanup
        server_task.abort();
        let _ = server_task.await;
    }

    /// Test that Abacus server starts and the whoami endpoint returns 401 Unauthorized
    #[test(sqlx::test)]
    async fn test_abacus_starts(pool: SqlitePool) {
        run_server_test(pool, |base_url| async move {
            let result = reqwest::get(format!("{base_url}/api/whoami"))
                .await
                .unwrap();

            assert_eq!(result.status(), 401);
        })
        .await;
    }

    /// Check all security headers are present with expected values
    #[test(sqlx::test)]
    async fn test_security_headers(pool: SqlitePool) {
        run_server_test(pool, |base_url| async move {
            let response = reqwest::get(format!("{base_url}/api/whoami"))
                .await
                .unwrap();

            assert_eq!(response.headers()["x-frame-options"], "deny");
            assert_eq!(response.headers()["x-content-type-options"], "nosniff");
            assert_eq!(
                response.headers()["x-permitted-cross-domain-policies"],
                "none"
            );
            assert_eq!(
                response.headers()["cross-origin-resource-policy"],
                "same-origin"
            );
            assert_eq!(
                response.headers()["cross-origin-embedder-policy"],
                "require-corp"
            );
            assert_eq!(
                response.headers()["cross-origin-opener-policy"],
                "same-origin"
            );

            #[cfg(feature = "memory-serve")]
            {
                let response = reqwest::get(format!("{base_url}/")).await.unwrap();
                assert_eq!(response.headers()["referrer-policy"], "no-referrer");
                assert_eq!(
                    response.headers()["permissions-policy"],
                    "accelerometer=(), autoplay=(), camera=(), cross-origin-isolated=(), \
                        display-capture=(), encrypted-media=(), fullscreen=(), geolocation=(), \
                        gyroscope=(), keyboard-map=(), magnetometer=(), microphone=(), midi=(), \
                        payment=(), picture-in-picture=(), publickey-credentials-get=(), \
                        screen-wake-lock=(), sync-xhr=(self), usb=(), web-share=(), \
                        xr-spatial-tracking=(), clipboard-read=(), clipboard-write=(), gamepad=(), \
                        hid=(), idle-detection=(), interest-cohort=(), serial=(), unload=()"
                );
                assert_eq!(
                    response.headers()["content-security-policy"],
                    "default-src 'self'; img-src 'self' data:"
                );
            }

            #[cfg(feature = "storybook")]
            {
                // Test that /storybook path doesn't have CSP header
                let storybook_response = reqwest::get(format!("{base_url}/storybook/"))
                    .await
                    .unwrap();
                assert!(
                    !storybook_response
                        .headers()
                        .contains_key("content-security-policy")
                );
            }
        })
        .await;
    }

    /// Check that the Cache-Control header is present with expected value
    #[test(sqlx::test)]
    async fn test_cache_headers(pool: SqlitePool) {
        run_server_test(pool, |base_url| async move {
            let response = reqwest::get(format!("{base_url}/api/whoami"))
                .await
                .unwrap();
            assert_eq!(response.headers()["cache-control"], "no-store");

            #[cfg(feature = "memory-serve")]
            {
                let response = reqwest::get(format!("{base_url}/")).await.unwrap();
                assert_eq!(response.headers()["cache-control"], "max-age:300, private");
            }
        })
        .await;
    }

    /// Check that unknown API routes return 404
    /// Note that this test is mostly designed for using with the memory-serve feature enabled.
    #[test(sqlx::test)]
    async fn test_api_fallback_route(pool: SqlitePool) {
        run_server_test(pool, |base_url| async move {
            // Test non-existent API endpoints
            let response = reqwest::get(format!("{base_url}/api")).await.unwrap();
            assert_eq!(response.status(), 404);

            let response = reqwest::get(format!("{base_url}/api/")).await.unwrap();
            assert_eq!(response.status(), 404);

            let response = reqwest::get(format!("{base_url}/api/nonexistent"))
                .await
                .unwrap();
            assert_eq!(response.status(), 404);

            let response = reqwest::get(format!("{base_url}/api/some/nested/path"))
                .await
                .unwrap();
            assert_eq!(response.status(), 404);

            // Test that valid API endpoints still work
            let response = reqwest::get(format!("{base_url}/api/whoami"))
                .await
                .unwrap();
            // Should return 401 (Unauthorized) since we're not logged in
            assert_eq!(response.status(), 401);
        })
        .await;
    }

    #[test(tokio::test)]
    async fn test_readonly_database_error() {
        // Create a read-only in-memory database
        let db_url = "sqlite::memory:?mode=ro";
        let result = create_sqlite_pool(
            db_url,
            #[cfg(feature = "dev-database")]
            false,
            #[cfg(feature = "dev-database")]
            false,
        )
        .await;

        assert!(matches!(result, Err(AppError::DatabaseReadOnly(_))));
    }
}
