// 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,
    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,
}

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" => (
            "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" => ("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()),
    };
    tokio::fs::read(app.root.join(path)).await.map_or_else(
        |_| secured(StatusCode::NOT_FOUND.into_response()),
        |bytes| secured(([("content-type", mime)], bytes).into_response()),
    )
}

pub(crate) async fn serve(
    listener: tokio::net::TcpListener,
    service: Service,
    root: &Path,
) -> 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(),
    };
    let routes = Router::new()
        .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(())
}