iocaine 2.0.0

// SPDX-FileCopyrightText: 2025 Gergely Nagy
// SPDX-FileContributor: Gergely Nagy
//
// SPDX-License-Identifier: MIT

use anyhow::Result;
use axum::{
    extract::{Path, Query, State},
    http::{header, HeaderMap, StatusCode},
    response::{IntoResponse, Response},
    routing::get,
    Router,
};
use std::collections::BTreeMap;
use std::sync::Arc;

use crate::{
    assembled_statistical_sequences::AssembledStatisticalSequences,
    config::Config,
    tenx_programmer::{TenXProgrammer, TenXProgrammerCounters},
};

#[derive(Debug, Clone)]
pub struct IocaineState {
    pub config: Config,
    pub counters: Option<TenXProgrammerCounters>,
    pub template: Arc<AssembledStatisticalSequences>,
}

#[derive(Debug)]
pub struct Iocaine {
    pub config: Config,
}

impl Iocaine {
    pub fn new(config: Config) -> Result<Self> {
        Ok(Self { config })
    }

    fn main_app(self, counters: Option<TenXProgrammerCounters>) -> Router {
        let state = IocaineState {
            config: self.config.clone(),
            template: Arc::new(AssembledStatisticalSequences::new(&self.config)),
            counters,
        };

        Router::new()
            .route("/", get(poison))
            .route("/{*path}", get(poison))
            .layer(tower_http::trace::TraceLayer::new_for_http())
            .with_state(state)
    }

    async fn start_server(self) -> std::result::Result<(), std::io::Error> {
        let bind = &self.config.server.bind.clone();
        let metrics_bind = self.config.metrics.bind.clone();
        let metrics_enabled = self.config.metrics.enable;

        let metrics = TenXProgrammer::new(&self.config.metrics);

        let app = self.main_app(metrics.as_ref().map(|v| v.counters.clone()));
        let listener = tokio::net::TcpListener::bind(bind).await?;
        let main_server = axum::serve(listener, app).with_graceful_shutdown(shutdown_signal());

        if metrics_enabled {
            let metrics_listener = tokio::net::TcpListener::bind(metrics_bind).await?;
            let metrics_app = metrics.unwrap().app();

            let metrics_server = axum::serve(metrics_listener, metrics_app)
                .with_graceful_shutdown(shutdown_signal());

            let _ = tokio::join!(main_server, metrics_server);

            Ok(())
        } else {
            Ok(main_server.await?)
        }
    }

    pub async fn run(self) -> Result<()> {
        self.start_server().await.unwrap();
        Ok(())
    }
}

pub async fn shutdown_signal() {
    let ctrl_c = async {
        tokio::signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

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

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

async fn poison(
    headers: axum::http::HeaderMap,
    State(state): State<IocaineState>,
    path: Option<Path<String>>,
    Query(params): Query<BTreeMap<String, String>>,
) -> std::result::Result<impl IntoResponse, AppError> {
    let default_host = axum::http::HeaderValue::from_static("<unknown>");
    let host = headers.get("host").unwrap_or(&default_host).to_str()?;
    let path = path.unwrap_or(Path("".to_string()));

    let (content_type, garbage) = state.template.generate(host, &path, &params)?;

    if let Some(counters) = state.counters {
        let labels = TenXProgrammer::build_label_values(&state.template.config.metrics, &headers)?;

        counters.request_counter.with_label_values(&labels).inc();
        counters
            .garbage_served_counter
            .with_label_values(&labels)
            .inc_by(garbage.len() as u64);

        let depth = path.chars().filter(|c| *c == '/').count() as u64;
        let maze_depth_counter = counters.maze_depth.with_label_values(&labels);
        let maze_depth = maze_depth_counter.get();
        if depth > maze_depth {
            maze_depth_counter.inc_by(depth - maze_depth);
        }
    }

    let mut headers = HeaderMap::new();
    headers.insert(header::CONTENT_TYPE, content_type.parse().unwrap());

    if state.config.templates.minify.enable
        && (content_type.starts_with("text/html") || content_type.starts_with("text/css"))
    {
        let config = &state.config.templates.minify;
        let cfg = minify_html::Cfg {
            minify_css: config.minify_css,
            minify_js: false,
            minify_doctype: false,
            ..Default::default()
        };
        let minified = minify_html::minify(garbage.as_bytes(), &cfg);
        Ok((headers, minified))
    } else {
        Ok((headers, garbage.into()))
    }
}

pub struct AppError(anyhow::Error);

impl IntoResponse for AppError {
    fn into_response(self) -> Response {
        tracing::error!("Internal server error: {}", self.0);
        (StatusCode::INTERNAL_SERVER_ERROR, "Something went wrong").into_response()
    }
}

impl From<axum::http::header::ToStrError> for AppError {
    fn from(e: axum::http::header::ToStrError) -> Self {
        Self(e.into())
    }
}

impl From<anyhow::Error> for AppError {
    fn from(e: anyhow::Error) -> Self {
        Self(e)
    }
}

impl From<std::io::Error> for AppError {
    fn from(e: std::io::Error) -> Self {
        Self(e.into())
    }
}

impl From<std::string::FromUtf8Error> for AppError {
    fn from(e: std::string::FromUtf8Error) -> Self {
        Self(e.into())
    }
}