// Copyright (C) 2026 Doan Luu Duc Tai <doan@luuduct.ai>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use crate::service::Service;
use anyhow::Result;
use axum::{
    Json, Router,
    body::Body,
    extract::{DefaultBodyLimit, State},
    http::{HeaderMap, HeaderValue, StatusCode},
    response::{IntoResponse, Response},
    routing::{get, post},
};
use serde_json::{Value, json};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use tokio::sync::Semaphore;

#[derive(Clone)]
struct App {
    service: Arc<Mutex<Service>>,
    permits: Arc<Semaphore>,
    root: PathBuf,
    origin: String,
    public: std::collections::BTreeMap<String, String>,
}

fn secured(mut response: Response) -> Response {
    for (name, value) in [
        ("cache-control", "no-store, max-age=0"),
        (
            "content-security-policy",
            "default-src 'none'; script-src 'self' 'wasm-unsafe-eval'; worker-src 'self'; connect-src 'self'; style-src 'self'; img-src 'none'; base-uri 'none'; form-action 'none'; frame-ancestors 'none'; object-src 'none'",
        ),
        ("referrer-policy", "no-referrer"),
        ("x-content-type-options", "nosniff"),
        ("cross-origin-opener-policy", "same-origin"),
        ("cross-origin-resource-policy", "same-origin"),
        (
            "permissions-policy",
            "camera=(), microphone=(), geolocation=(), payment=(), usb=()",
        ),
    ] {
        response
            .headers_mut()
            .insert(name, HeaderValue::from_static(value));
    }
    response
}

async fn api(State(app): State<App>, headers: HeaderMap, Json(request): Json<Value>) -> Response {
    if headers.get("origin").and_then(|value| value.to_str().ok()) != Some(app.origin.as_str()) {
        return secured(
            (
                StatusCode::FORBIDDEN,
                Json(json!({"error":"Origin rejected"})),
            )
                .into_response(),
        );
    }
    let Ok(permit) = Arc::clone(&app.permits).try_acquire_owned() else {
        return secured(
            (
                StatusCode::TOO_MANY_REQUESTS,
                Json(json!({"error":"Server busy"})),
            )
                .into_response(),
        );
    };
    let result = tokio::task::spawn_blocking(move || {
        let _permit = permit;
        app.service
            .lock()
            .map_err(|_| anyhow::anyhow!("Service unavailable"))?
            .call(&request)
    })
    .await;
    match result {
        Ok(Ok(value)) => secured(Json(value).into_response()),
        _ => secured((StatusCode::BAD_REQUEST,Json(json!({"error":"Request rejected. Unlock again if your session expired; check credentials, message limits, and conversation access."}))).into_response()),
    }
}

async fn asset(State(app): State<App>, uri: axum::http::Uri) -> Response {
    let (path, mime) = match uri.path() {
        "/" => (
            "crates/chat-server/web/index.html",
            "text/html; charset=utf-8",
        ),
        "/app.js" => (
            "crates/chat-server/web/app.js",
            "text/javascript; charset=utf-8",
        ),
        "/worker.js" => (
            "crates/chat-server/web/worker.js",
            "text/javascript; charset=utf-8",
        ),
        "/wasi.mjs" => (
            "crates/chat-crypto/browser/wasi.mjs",
            "text/javascript; charset=utf-8",
        ),
        "/style.css" => (
            "crates/chat-server/web/style.css",
            "text/css; charset=utf-8",
        ),
        "/crypto.wasm" => (
            if app.root.join("bin/chat/crypto.wasm").is_file() {
                "bin/chat/crypto.wasm"
            } else {
                "target/wasm32-wasip1/release/stack-chat-crypto.wasm"
            },
            "application/wasm",
        ),
        "/javascript.html" => (
            "crates/chat-server/web/javascript.html",
            "text/html; charset=utf-8",
        ),
        "/SOURCE.tar.gz" => (
            if app.root.join("bin/chat/SOURCE.tar.gz").is_file() {
                "bin/chat/SOURCE.tar.gz"
            } else {
                "var/build/chat-source.tar.gz"
            },
            "application/gzip",
        ),
        "/THIRD-PARTY-LICENSES.html" => {
            ("web/THIRD-PARTY-LICENSES.html", "text/html; charset=utf-8")
        }
        "/LIMITATIONS.txt" => ("docs/CHAT-APPLICATION.txt", "text/plain; charset=utf-8"),
        "/LICENSE" => ("LICENSE", "text/plain; charset=utf-8"),
        _ => return secured(StatusCode::NOT_FOUND.into_response()),
    };
    if matches!(uri.path(), "/" | "/javascript.html") {
        let Ok(bytes) = tokio::fs::read(app.root.join(path)).await else {
            return secured(StatusCode::NOT_FOUND.into_response());
        };
        let rendered = String::from_utf8(bytes)
            .map_err(anyhow::Error::from)
            .and_then(|source| stack_config::render(&source, &app.public));
        return rendered.map_or_else(
            |_| secured(StatusCode::INTERNAL_SERVER_ERROR.into_response()),
            |html| secured(([("content-type", mime)], html).into_response()),
        );
    }
    let Ok(file) = tokio::fs::File::open(app.root.join(path)).await else {
        return secured(StatusCode::NOT_FOUND.into_response());
    };
    let length = file.metadata().await.ok().map(|metadata| metadata.len());
    let mut response = Response::new(Body::from_stream(tokio_util::io::ReaderStream::new(file)));
    response
        .headers_mut()
        .insert("content-type", HeaderValue::from_static(mime));
    if let Some(header_length) =
        length.and_then(|value| HeaderValue::from_str(&value.to_string()).ok())
    {
        response
            .headers_mut()
            .insert("content-length", header_length);
    }
    secured(response)
}

pub(crate) async fn serve(
    listener: tokio::net::TcpListener,
    service: Service,
    root: &Path,
    public: std::collections::BTreeMap<String, String>,
) -> Result<()> {
    let app = App {
        origin: service.origin.clone(),
        service: Arc::new(Mutex::new(service)),
        permits: Arc::new(Semaphore::new(4)),
        root: root.to_path_buf(),
        public,
    };
    let routes = Router::new()
        .route(
            "/healthz",
            get(|| async { secured("ok\n".into_response()) }),
        )
        .route("/api", post(api))
        .fallback(get(asset))
        .layer(DefaultBodyLimit::max(262_144))
        .with_state(app);
    axum::serve(listener, routes)
        .with_graceful_shutdown(async {
            drop(tokio::signal::ctrl_c().await);
        })
        .await?;
    Ok(())
}