use log::info;
use once_cell::sync::Lazy;
use axum::extract::DefaultBodyLimit;
use axum::middleware;
use axum::routing::{delete, get, post};
use axum::Router;
use axum_server::tls_rustls::RustlsConfig;
use std::net::SocketAddr;
use std::sync::Once;
use crate::api_token::APIToken;
use crate::runtime_certificate::{CERTIFICATE, CERTIFICATE_PRIVATE_KEY};
use crate::environment::is_dev;
use crate::network::get_available_port;

static RUSTLS_CRYPTO_PROVIDER_INIT: Once = Once::new();

/// The request body limit for one batch of prompt injection filtering. The app caps the text
/// it returns to a model well below this, so the limit is headroom, not a working constraint.
const PROMPT_INJECTION_BATCH_BODY_LIMIT_BYTES: usize = 16 * 1024 * 1024;

/// The port used for the runtime API server. In the development environment, we use a fixed
/// port, in the production environment we use the next available port. This differentiation
/// is necessary because we cannot communicate the port to the .NET server in the development
/// environment.
pub static API_SERVER_PORT: Lazy<u16> = Lazy::new(|| {
    if is_dev() {
        5000
    } else {
        get_available_port().unwrap()
    }
});

/// Creates the router with every route of the runtime API, each of which requires the API token.
fn create_router() -> Router {
    let router = Router::new()
        .route("/system/dotnet/port", get(crate::dotnet::dotnet_port))
        .route("/system/dotnet/ready", get(crate::dotnet::dotnet_ready))
        .route("/system/qdrant-edge/info", get(crate::qdrant_edge_database::qdrant_edge_info))
        .route("/system/qdrant-edge/ensure", post(crate::qdrant_edge_database::ensure_qdrant_edge_store))
        .route("/system/qdrant-edge/exists", post(crate::qdrant_edge_database::check_qdrant_edge_store))
        .route("/system/qdrant-edge/insert", post(crate::qdrant_edge_database::insert_qdrant_edge_embedding))
        .route("/system/qdrant-edge/search", post(crate::qdrant_edge_database::search_qdrant_edge_embeddings))
        .route("/system/qdrant-edge/delete-file", post(crate::qdrant_edge_database::delete_qdrant_edge_embedding_by_file))
        .route("/system/qdrant-edge/optimize", post(crate::qdrant_edge_database::optimize_qdrant_edge_store))
        .route("/system/qdrant-edge/delete-store", post(crate::qdrant_edge_database::delete_qdrant_edge_store))
        .route("/system/qdrant-edge/backup/prepare", post(crate::qdrant_edge_database::prepare_qdrant_edge_backup))
        .route("/system/qdrant-edge/backup/copy", post(crate::qdrant_edge_database::copy_qdrant_edge_store_for_backup))
        .route("/clipboard/set", post(crate::clipboard::set_clipboard))
        .route("/share/file", post(crate::share_sheet::share_file))
        .route("/events", get(crate::app_window::get_event_stream))
        .route("/updates/check", get(crate::app_window::check_for_update))
        .route("/updates/install", get(crate::app_window::install_update))
        .route("/app/exit", post(crate::app_window::exit_app))
        .route("/app/restart", post(crate::app_window::restart_app))
        .route("/select/directory", post(crate::file_actions::select_directory))
        .route("/select/file", post(crate::file_actions::select_file))
        .route("/select/files", post(crate::file_actions::select_files))
        .route("/save/file", post(crate::file_actions::save_file))
        .route("/open/path", post(crate::file_actions::open_path_in_file_manager))
        .route("/open/document", post(crate::file_actions::open_document))
        .route("/secrets/get", post(crate::secret::get_secret))
        .route("/secrets/store", post(crate::secret::store_secret))
        .route("/secrets/delete", post(crate::secret::delete_secret))
        .route("/system/directories/config", get(crate::environment::get_config_directory))
        .route("/system/directories/data", get(crate::environment::get_data_directory))
        .route("/system/runtime/info", get(crate::environment::get_runtime_info))
        .route("/system/language", get(crate::environment::read_user_language))
        .route("/system/username", get(crate::environment::read_user_name))
        .route("/system/enterprise/config/encryption_secret", get(crate::environment::read_enterprise_env_config_encryption_secret))
        .route("/system/enterprise/configs", get(crate::environment::read_enterprise_configs))
        .route("/retrieval/fs/extract", get(crate::file_data::extract_data))
        .route("/security/prompt-injection/sanitize", post(crate::prompt_injection::api::sanitize))
        //
        // A batch carries every text of one tool call, which is far more than Axum's 2 MB
        // default allows. Exceeding that limit would answer 413, and the app treats a failed
        // filter call as "cannot filter" and uses the text unfiltered — the protection would
        // drop out silently on exactly the largest results. Hence the explicit limit.
        //
        .route("/security/prompt-injection/sanitize-batch", post(crate::prompt_injection::api::sanitize_batch)
            .layer(DefaultBodyLimit::max(PROMPT_INJECTION_BATCH_BODY_LIMIT_BYTES)))
        .route("/media/jobs", post(crate::media::create_job))
        .route("/media/jobs/{id}/events", get(crate::media::get_job_events))
        .route("/media/jobs/{id}", delete(crate::media::cancel_job))
        .route("/image/prepare", post(crate::image::prepare_image))
        .route("/log/paths", get(crate::log::get_log_paths))
        .route("/log/event", post(crate::log::log_event))
        .route("/tokenizer/count", post(crate::tokenizer::token_count))
        .route("/tokenizer/validate", post(crate::tokenizer::validate_tokenizer))
        .route("/tokenizer/store", post(crate::tokenizer::store_tokenizer))
        .route("/tokenizer/delete", post(crate::tokenizer::delete_tokenizer))
        .route("/shortcuts/register", post(crate::app_window::register_shortcut))
        .route("/shortcuts/validate", post(crate::app_window::validate_shortcut))
        .route("/shortcuts/suspend", post(crate::app_window::suspend_shortcuts))
        .route("/shortcuts/resume", post(crate::app_window::resume_shortcuts));

    require_api_token(router)
}

/// Requires a valid API token for every route of the given router.
///
/// This layer is the only place where the runtime API checks the token, so the handlers take no
/// `APIToken` argument. Asking each handler for it left a route open to every process on this
/// machine as soon as one handler forgot it.
///
/// `route_layer` protects only the routes registered before it, so call this function after the
/// last route. Unknown paths still answer 404 without checking the token.
fn require_api_token(router: Router) -> Router {
    router.route_layer(middleware::from_extractor::<APIToken>())
}

/// Starts the runtime API server. The server is used to communicate with the .NET server and
/// to provide additional functionality to the Tauri app.
pub fn start_runtime_api() {
    let api_port = *API_SERVER_PORT;
    info!("Try to start the API server on 'http://localhost:{api_port}'...");

    let app = create_router();

    tauri::async_runtime::spawn(async move {
        install_rustls_crypto_provider();

        let cert = CERTIFICATE.get().unwrap().clone();
        let key = CERTIFICATE_PRIVATE_KEY.get().unwrap().clone();
        let tls_config = RustlsConfig::from_pem(cert, key).await.unwrap();
        let addr = SocketAddr::from(([127, 0, 0, 1], api_port));

        axum_server::bind_rustls(addr, tls_config)
            .serve(app.into_make_service())
            .await
            .unwrap();
    });
}

fn install_rustls_crypto_provider() {
    RUSTLS_CRYPTO_PROVIDER_INIT.call_once(|| {
        let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
    });
}

#[cfg(test)]
pub(crate) mod test_support {
    use std::sync::Mutex;
    use std::time::{Duration, Instant};

    /// How long the lock stays held while the request waits for it.
    const LOCK_DURATION: Duration = Duration::from_secs(1);

    /// How long the probe may wait for the runtime. Far below the lock duration, so a request which
    /// blocks the worker while it waits fails the test every time.
    const PROBE_LIMIT: Duration = Duration::from_millis(200);

    /// Starts a request while another thread holds the given lock, as if a long operation were
    /// running, and checks that a probe still gets through the runtime meanwhile. Returns the
    /// response once the lock is free again.
    ///
    /// Call it from a test with a single worker, `#[tokio::test(flavor = "multi_thread",
    /// worker_threads = 1)]`. A request which blocks that worker while it waits then holds up the
    /// probe, just as in the app it holds up the one HTTP/2 connection which all calls share.
    pub(crate) async fn assert_runtime_stays_free<T, F>(lock: &'static Mutex<T>, request: F) -> F::Output
    where
        T: Send + 'static,
        F: Future + Send + 'static,
        F::Output: Send + 'static,
    {
        let (locked_sender, locked_receiver) = std::sync::mpsc::channel();
        let lock_holder = std::thread::spawn(move || {
            let guard = lock.lock().unwrap();
            locked_sender.send(()).unwrap();
            std::thread::sleep(LOCK_DURATION);
            drop(guard);
        });

        locked_receiver.recv().unwrap();
        let request = tokio::spawn(request);

        let probe_started = Instant::now();
        tokio::spawn(async {}).await.unwrap();
        let probe_duration = probe_started.elapsed();

        lock_holder.join().unwrap();
        let response = request.await.unwrap();

        assert!(probe_duration < PROBE_LIMIT, "The probe waited {probe_duration:?} for the runtime.");
        response
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;
    use axum::body::Body;
    use axum::http::{Request, StatusCode};
    use regex::Regex;
    use tokio::time::timeout;
    use tower::ServiceExt;
    use crate::runtime_api_token::API_TOKEN;

    /// How long a route may take to refuse a request without the token. The refusal needs no work,
    /// so a route which takes longer has passed the request on to its handler. In a test, some
    /// handlers then never answer, for example one which waits for the main window, and others
    /// panic because the app around them is missing.
    const REFUSAL_LIMIT: Duration = Duration::from_secs(1);

    /// One route which `create_router` registers.
    struct RegisteredRoute {
        method: String,
        path: String,
    }

    /// Reads every route from the source of `create_router`, so that a new route is checked
    /// without anybody having to add it to a test.
    fn registered_routes() -> Vec<RegisteredRoute> {
        let source = include_str!("runtime_api.rs");
        let start = source.find("fn create_router()").unwrap();
        let end = start + source[start..].find("\n}").unwrap();
        let create_router_source = &source[start..end];

        let route_pattern = Regex::new(r#"\.route\("([^"]+)",\s*(get|post|put|patch|delete)\("#).unwrap();
        let path_parameter_pattern = Regex::new(r"\{[^}]+\}").unwrap();
        let routes: Vec<RegisteredRoute> = route_pattern
            .captures_iter(create_router_source)
            .map(|route| RegisteredRoute {
                method: route[2].to_uppercase(),
                path: path_parameter_pattern.replace_all(&route[1], "x").into_owned(),
            })
            .collect();

        //
        // A route which the pattern cannot read would silently drop out of the check, for example
        // one whose methods are not a plain get, post, put, patch, or delete call. Extend the
        // pattern then.
        //
        assert_eq!(
            routes.len(),
            create_router_source.matches(".route(").count(),
            "The test could not read every route of create_router.",
        );

        assert!(!routes.is_empty(), "The test found no route in create_router.");
        routes
    }

    /// Sends one request through the router, without a server, and returns the status of the response.
    async fn send(router: Router, request: Request<Body>) -> StatusCode {
        router.oneshot(request).await.unwrap().status()
    }

    /// Builds a request with the given token header, or without one.
    fn request(method: &str, uri: &str, token: Option<&str>, body: &'static str) -> Request<Body> {
        let mut builder = Request::builder()
            .method(method)
            .uri(uri)
            .header("content-type", "application/json");

        if let Some(token) = token {
            builder = builder.header("token", token);
        }

        builder.body(Body::from(body)).unwrap()
    }

    #[tokio::test]
    async fn require_api_token_lets_only_the_valid_token_through() {
        let router = require_api_token(Router::new().route("/probe", get(|| async {})));

        assert_eq!(send(router.clone(), request("GET", "/probe", None, "")).await, StatusCode::UNAUTHORIZED);
        assert_eq!(send(router.clone(), request("GET", "/probe", Some("0123abcd"), "")).await, StatusCode::UNAUTHORIZED);
        assert_eq!(send(router, request("GET", "/probe", Some(API_TOKEN.to_hex_text()), "")).await, StatusCode::OK);
    }

    #[tokio::test]
    async fn every_route_requires_the_api_token() {
        let mut unprotected_routes = Vec::new();
        for route in registered_routes() {
            // A task of its own turns a panicking handler into a finding instead of ending the test.
            let answer = tokio::spawn(send(create_router(), request(&route.method, &route.path, None, "")));
            match timeout(REFUSAL_LIMIT, answer).await {
                Ok(Ok(StatusCode::UNAUTHORIZED)) => {},
                Ok(Ok(status)) => unprotected_routes.push(format!("{} {} answered {status}", route.method, route.path)),
                Ok(Err(_)) => unprotected_routes.push(format!("{} {} panicked in its handler", route.method, route.path)),
                Err(_) => unprotected_routes.push(format!("{} {} did not refuse within {REFUSAL_LIMIT:?}", route.method, route.path)),
            }
        }

        assert!(unprotected_routes.is_empty(), "These routes answer without the API token: {unprotected_routes:#?}");
    }

    #[tokio::test]
    async fn a_valid_token_gets_through_to_the_handler() {
        // The handler fails on the incomplete body, before it touches any file.
        let with_token = request("POST", "/image/prepare", Some(API_TOKEN.to_hex_text()), "{}");
        assert_eq!(send(create_router(), with_token).await, StatusCode::UNPROCESSABLE_ENTITY);
    }
}