use axum::{
    extract::{RawQuery, State},
    http::{header, HeaderValue, StatusCode},
    middleware,
    response::{IntoResponse, Response},
    routing::get,
    Router,
};
use forma_ir::PageProps;
use forma_server::{assets, load_ir_modules, sw, PageHandle, RenderCtx, RouteMode};
use rust_embed::Embed;
use std::{net::IpAddr, path::PathBuf, sync::Arc};
use tower_http::services::{ServeDir, ServeFile};

mod paywall;

#[derive(Embed)]
#[folder = "dist/"]
struct Assets;

#[derive(PageProps)]
struct DemoProps {
    count: u32,
    raw_html_href: String,
}

struct AppState {
    islands: PageHandle<DemoProps>,
    no_js: PageHandle<DemoProps>,
    paywall: paywall::Paywall,
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    tracing_subscriber::fmt::init();
    let manifest = assets::try_load_manifest::<Assets>()?;
    for path in ["/lab/native/", "/lab/native/no-js/", "/lab/paywall/"] {
        if manifest.route(path).map(|route| route.mode) != Some(RouteMode::Ssr) {
            return Err(
                format!("{path} requires a native SSR build; run node native/build.mjs").into(),
            );
        }
    }
    for path in ["/lab/native/no-js/", "/lab/paywall/"] {
        if !manifest.route(path).unwrap().js.is_empty() {
            return Err(format!("{path} must not ship JavaScript").into());
        }
    }
    let pages = load_ir_modules::<Assets>(&manifest).map_err(|errors| format!("{errors:?}"))?;
    let islands = pages
        .bind::<DemoProps>("islands")
        .map_err(|errors| format!("{errors:?}"))?;
    let no_js = pages
        .bind::<DemoProps>("no-js")
        .map_err(|errors| format!("{errors:?}"))?;
    let paywall_page = pages
        .bind::<paywall::PaywallProps>("paywall")
        .map_err(|errors| format!("{errors:?}"))?;
    pages.finalize().map_err(|errors| format!("{errors:?}"))?;
    sw::configure(sw::ServiceWorker::off()).map_err(|error| format!("{error:?}"))?;

    let site_dist = std::env::var_os("SITE_DIST")
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../dist"));
    if !site_dist.join("index.html").is_file() {
        return Err(format!(
            "SITE_DIST {} has no index.html; build the website first",
            site_dist.display()
        )
        .into());
    }
    let paywall = paywall::Paywall::load(paywall_page, &site_dist)?;
    let state = Arc::new(AppState {
        islands,
        no_js,
        paywall,
    });
    // Portal documents, downloads and browser tools use mutable filenames.
    // Revalidate these responses on every reuse instead of allowing heuristic
    // freshness to mix old playground code with a new deployment.
    let static_site = Router::<()>::new()
        .fallback_service(
            ServeDir::new(&site_dist)
                .append_index_html_on_directories(true)
                .not_found_service(ServeFile::new(site_dist.join("404.html"))),
        )
        .layer(middleware::map_response(revalidate_static));
    let app = Router::new()
        .merge(paywall::routes())
        .route(
            "/favicon.ico",
            get(|| async { axum::response::Redirect::permanent("/favicon.svg") }),
        )
        .route("/lab/native", get(island_demo))
        .route("/lab/native/", get(island_demo))
        .route("/lab/native/html/", get(island_html))
        .route(
            "/lab/native/inspect/",
            get(|| async { inspection("index.html", "text/html; charset=utf-8") }),
        )
        .route(
            "/lab/native/inspect/report.json",
            get(|| async { inspection("report.json", "application/json") }),
        )
        .route(
            "/lab/native/inspect/manifest.json",
            get(|| async { inspection("manifest.json", "application/json") }),
        )
        .route(
            "/lab/native/inspect/islands.txt",
            get(|| async { inspection("islands.txt", "text/plain; charset=utf-8") }),
        )
        .route(
            "/lab/native/inspect/no-js.txt",
            get(|| async { inspection("no-js.txt", "text/plain; charset=utf-8") }),
        )
        .route(
            "/lab/native/inspect/islands.ir",
            get(|| async { inspection("islands.ir", "application/octet-stream") }),
        )
        .route(
            "/lab/native/inspect/no-js.ir",
            get(|| async { inspection("no-js.ir", "application/octet-stream") }),
        )
        .route(
            "/lab/native/source/islands.tsx",
            get(|| async { source(include_str!("islands.tsx")) }),
        )
        .route(
            "/lab/native/source/no-js.tsx",
            get(|| async { source(include_str!("no-js.tsx")) }),
        )
        .route(
            "/lab/native/source/server.rs",
            get(|| async { source(include_str!("main.rs")) }),
        )
        .route("/lab/native/no-js", get(no_js_demo))
        .route("/lab/native/no-js/", get(no_js_demo))
        .route("/lab/native/no-js/html/", get(no_js_html))
        .route("/healthz", get(|| async { "ok\n" }))
        .route("/sw.js", get(sw::serve_sw::<Assets>))
        .route("/_assets/{*filename}", get(assets::serve_asset::<Assets>))
        .fallback_service(static_site)
        .with_state(state);
    let port = std::env::var("PORT")
        .unwrap_or_else(|_| "4174".to_owned())
        .parse::<u16>()?;
    let host = std::env::var("HOST")
        .unwrap_or_else(|_| "127.0.0.1".to_owned())
        .parse::<IpAddr>()?;
    let listener = tokio::net::TcpListener::bind((host, port)).await?;
    println!("FORMA_NATIVE_LISTENING http://{}", listener.local_addr()?);
    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal())
        .await?;
    Ok(())
}

async fn revalidate_static(mut response: Response) -> Response {
    response
        .headers_mut()
        .insert(header::CACHE_CONTROL, HeaderValue::from_static("no-cache"));
    response
}

async fn shutdown_signal() {
    let interrupt = async {
        tokio::signal::ctrl_c()
            .await
            .expect("failed to install the Ctrl-C handler");
    };
    // Container hosts stop the process with SIGTERM, while local development
    // normally uses Ctrl-C. Either signal must let Axum drain active responses.
    #[cfg(unix)]
    let terminate = async {
        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
            .expect("failed to install the SIGTERM handler")
            .recv()
            .await;
    };
    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = interrupt => {},
        _ = terminate => {},
    }
}

fn request_count(query: Option<&str>) -> Result<u32, &'static str> {
    let Some(query) = query.filter(|query| !query.is_empty()) else {
        return Ok(7);
    };
    if query.len() > 256 {
        return Err("Query is too long. Use one count from 0 to 9999.");
    }
    let pairs = serde_urlencoded::from_str::<Vec<(String, String)>>(query)
        .map_err(|_| "Invalid query. Use one count from 0 to 9999.")?;
    if pairs.len() != 1 || pairs[0].0 != "count" {
        return Err("Use exactly one count from 0 to 9999.");
    }
    let value = &pairs[0].1;
    if value.is_empty() || value.len() > 4 || !value.bytes().all(|byte| byte.is_ascii_digit()) {
        return Err("Count must be an integer from 0 to 9999.");
    }
    value
        .parse::<u32>()
        .map_err(|_| "Count must be an integer from 0 to 9999.")
}

fn render_demo(
    page: &PageHandle<DemoProps>,
    query: Option<&str>,
    title: &str,
    raw_path: &str,
) -> Response {
    let count = match request_count(query) {
        Ok(count) => count,
        Err(message) => {
            return (
                StatusCode::BAD_REQUEST,
                [
                    (header::CACHE_CONTROL, "no-store"),
                    (header::X_CONTENT_TYPE_OPTIONS, "nosniff"),
                ],
                message,
            )
                .into_response()
        }
    };
    let props = DemoProps {
        count,
        raw_html_href: format!("{raw_path}?count={count}"),
    };
    match page.render(&props, &RenderCtx::new().with_title(title)) {
        // Keep each response's HTML, fresh nonce, CSP and no-store headers together.
        Ok(response) => response.into_response(),
        Err(error) => {
            tracing::error!(code = error.code(), route = error.route(), "{error}");
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                [
                    (header::CACHE_CONTROL, "no-store"),
                    (header::X_CONTENT_TYPE_OPTIONS, "nosniff"),
                ],
                "The native demonstration could not be rendered.",
            )
                .into_response()
        }
    }
}

async fn island_demo(State(state): State<Arc<AppState>>, RawQuery(query): RawQuery) -> Response {
    render_demo(
        &state.islands,
        query.as_deref(),
        "Native islands lab — Forma",
        "/lab/native/html/",
    )
}

async fn no_js_demo(State(state): State<Arc<AppState>>, RawQuery(query): RawQuery) -> Response {
    render_demo(
        &state.no_js,
        query.as_deref(),
        "Zero-JavaScript lab — Forma",
        "/lab/native/no-js/html/",
    )
}

fn as_plain_text(mut response: Response) -> Response {
    // Inspection view: retain the rendered response's no-store and security
    // headers; display its body as text without executing the document.
    response.headers_mut().insert(
        header::CONTENT_TYPE,
        "text/plain; charset=utf-8".parse().unwrap(),
    );
    response
}

async fn island_html(State(state): State<Arc<AppState>>, RawQuery(query): RawQuery) -> Response {
    as_plain_text(render_demo(
        &state.islands,
        query.as_deref(),
        "Native islands lab — Forma",
        "/lab/native/html/",
    ))
}

async fn no_js_html(State(state): State<Arc<AppState>>, RawQuery(query): RawQuery) -> Response {
    as_plain_text(render_demo(
        &state.no_js,
        query.as_deref(),
        "Zero-JavaScript lab — Forma",
        "/lab/native/no-js/html/",
    ))
}

fn source(text: &'static str) -> Response {
    (
        [
            (header::CONTENT_TYPE, "text/plain; charset=utf-8"),
            (header::CACHE_CONTROL, "no-store"),
            (header::X_CONTENT_TYPE_OPTIONS, "nosniff"),
        ],
        text,
    )
        .into_response()
}

fn inspection(name: &'static str, content_type: &'static str) -> Response {
    let Some(file) = Assets::get(&format!("inspection/{name}")) else {
        return StatusCode::NOT_FOUND.into_response();
    };
    (
        [
            (header::CONTENT_TYPE, content_type),
            (header::CACHE_CONTROL, "no-store"),
            (header::X_CONTENT_TYPE_OPTIONS, "nosniff"),
            // This is a static inspection document, not a rendered PageResponse.
            (header::CONTENT_SECURITY_POLICY, "default-src 'self'; script-src 'self' 'wasm-unsafe-eval'; style-src 'self'; img-src 'self' data:; font-src 'self'; object-src 'none'; base-uri 'none'; frame-ancestors 'none'; form-action 'self'"),
        ],
        file.data.into_owned(),
    ).into_response()
}

#[cfg(test)]
mod tests {
    use super::request_count;

    #[test]
    fn bounds_request_input_before_rendering() {
        assert_eq!(request_count(None), Ok(7));
        assert_eq!(request_count(Some("")), Ok(7));
        assert_eq!(request_count(Some("count=0")), Ok(0));
        assert_eq!(request_count(Some("count=9999")), Ok(9999));
        for query in [
            "count=10000",
            "count=-1",
            "count=1.5",
            "count=",
            "count=NaN",
            "count=%3Cscript%3E",
            "count=7&count=8",
            "other=7",
        ] {
            assert!(
                request_count(Some(query)).is_err(),
                "unexpectedly accepted {query}"
            );
        }
        assert!(request_count(Some(&"count=1&".repeat(40))).is_err());
    }
}