Skip to content
pub mod compress;
pub mod db;
pub mod models;

use actix_web::{App, Error, HttpRequest, HttpResponse, HttpServer, Responder, get, post, web};
use compress::{ContentEncoding, ContentType};
use std::env;

pub struct AppState {
    pub db_pool: sqlx::SqlitePool,
    pub db_path: String,
    pub max_db_size: u64,
    pub api_key: Option<String>,
}

#[get("/up")]
async fn up() -> impl Responder {
    HttpResponse::Ok().body("up")
}

#[post("/ingest")]
async fn ingest(
    req: HttpRequest,
    payload: web::Bytes,
    state: web::Data<AppState>,
) -> Result<HttpResponse, Error> {
    if let Some(ref expected_key) = state.api_key {
        let authenticated = req
            .headers()
            .get("Authorization")
            .and_then(|h| h.to_str().ok())
            .and_then(|s| s.strip_prefix("Bearer ").map(|token| token.trim()))
            .is_some_and(|key| key == expected_key);

        if !authenticated {
            return Ok(
                HttpResponse::Unauthorized().body("Unauthorized: Invalid or missing API key")
            );
        }
    }

    let encoding = req
        .headers()
        .get("X-Content-Encoding")
        .or_else(|| req.headers().get("Content-Encoding"))
        .and_then(|h| h.to_str().ok())
        .and_then(|s| s.parse::<ContentEncoding>().ok())
        .unwrap_or(ContentEncoding::Identity);

    let content_type = req
        .headers()
        .get("Content-Type")
        .and_then(|h| h.to_str().ok())
        .and_then(|s| s.parse::<ContentType>().ok())
        .unwrap_or(ContentType::Json);

    let logs = match compress::parse_payload(&payload, encoding, content_type) {
        Ok(l) => l,
        Err(e) => {
            return Ok(HttpResponse::BadRequest().body(format!("Failed to parse payload: {}", e)));
        }
    };

    let count = logs.len();

    if let Err(e) = db::insert_logs(&state.db_pool, &logs).await {
        return Ok(
            HttpResponse::InternalServerError().body(format!("Database insert error: {}", e))
        );
    }

    if let Err(e) = db::enforce_size_limit(&state.db_pool, &state.db_path, state.max_db_size).await
    {
        eprintln!("Failed to enforce size limit: {}", e);
    }

    Ok(HttpResponse::Ok().json(serde_json::json!({
        "status": "success",
        "ingested": count
    })))
}

#[actix_web::main]
async fn main() -> std::io::Result<()> {
    let db_path = env::var("DATABASE_PATH").unwrap_or_else(|_| "logs.db".to_string());
    let max_db_size: u64 = env::var("MAX_DB_SIZE_BYTES")
        .ok()
        .and_then(|s| s.parse().ok())
        .unwrap_or(1_073_741_824); // 1 GB in bytes

    println!("Initializing database at {}", db_path);
    let db_pool = db::init_db(&db_path).await.map_err(std::io::Error::other)?;

    let api_key = env::var("API_KEY").ok();

    let state = web::Data::new(AppState {
        db_pool,
        db_path,
        max_db_size,
        api_key,
    });

    println!("Starting server at 127.0.0.1:9180");
    HttpServer::new(move || {
        App::new()
            .app_data(state.clone())
            .service(up)
            .service(ingest)
    })
    .bind(("127.0.0.1", 9180))?
    .run()
    .await
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::models::LogLine;
    use actix_web::{App, test};
    use chrono::Utc;
    use flate2::Compression;
    use flate2::write::GzEncoder;
    use serde_json::json;
    use std::collections::HashMap;
    use std::io::Write;

    fn compress_zstd(data: &[u8]) -> Vec<u8> {
        let mut encoder = zstd::stream::write::Encoder::new(Vec::new(), 0).unwrap();
        encoder.write_all(data).unwrap();
        encoder.finish().unwrap()
    }

    fn compress_gzip(data: &[u8]) -> Vec<u8> {
        let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
        encoder.write_all(data).unwrap();
        encoder.finish().unwrap()
    }

    fn make_test_logs(count: usize) -> Vec<LogLine> {
        (0..count)
            .map(|i| LogLine {
                timestamp: Utc::now() - chrono::Duration::seconds((count - i) as i64),
                severity: if i % 2 == 0 {
                    "INFO".to_string()
                } else {
                    "ERROR".to_string()
                },
                message: format!("Log message {}", i),
                service: "test-service".to_string(),
                metadata: {
                    let mut m = HashMap::new();
                    m.insert("index".to_string(), json!(i));
                    m
                },
            })
            .collect()
    }

    #[actix_web::test]
    async fn test_up_endpoint() {
        let app = test::init_service(App::new().service(up)).await;
        let req = test::TestRequest::get().uri("/up").to_request();
        let resp = test::call_service(&app, req).await;
        assert!(resp.status().is_success());
        let body = test::read_body(resp).await;
        assert_eq!(body, "up");
    }

    #[actix_web::test]
    async fn test_ingest_json_uncompressed() {
        let db_path = "test_ingest_json_uncompressed.db";
        let _ = std::fs::remove_file(db_path);

        let db_pool = db::init_db(db_path).await.unwrap();
        let state = web::Data::new(AppState {
            db_pool: db_pool.clone(),
            db_path: db_path.to_string(),
            max_db_size: 1_000_000,
            api_key: None,
        });

        let app = test::init_service(App::new().app_data(state.clone()).service(ingest)).await;

        let logs = make_test_logs(5);
        let payload = serde_json::to_vec(&logs).unwrap();

        let req = test::TestRequest::post()
            .uri("/ingest")
            .insert_header(("Content-Type", "application/json"))
            .set_payload(payload)
            .to_request();

        let resp = test::call_service(&app, req).await;
        assert!(resp.status().is_success());

        // Verify count
        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM logs")
            .fetch_one(&db_pool)
            .await
            .unwrap();
        assert_eq!(row.0, 5);

        db_pool.close().await;
        let _ = std::fs::remove_file(db_path);
        let _ = std::fs::remove_file(format!("{}-wal", db_path));
        let _ = std::fs::remove_file(format!("{}-shm", db_path));
    }

    #[actix_web::test]
    async fn test_ingest_zstd_msgpack() {
        let db_path = "test_ingest_zstd_msgpack.db";
        let _ = std::fs::remove_file(db_path);

        let db_pool = db::init_db(db_path).await.unwrap();
        let state = web::Data::new(AppState {
            db_pool: db_pool.clone(),
            db_path: db_path.to_string(),
            max_db_size: 1_000_000,
            api_key: None,
        });

        let app = test::init_service(App::new().app_data(state.clone()).service(ingest)).await;

        let logs = make_test_logs(10);
        let msgpack_bytes = rmp_serde::to_vec(&logs).unwrap();
        let zstd_payload = compress_zstd(&msgpack_bytes);

        let req = test::TestRequest::post()
            .uri("/ingest")
            .insert_header(("Content-Type", "application/msgpack"))
            .insert_header(("X-Content-Encoding", "zstd"))
            .set_payload(zstd_payload)
            .to_request();

        let resp = test::call_service(&app, req).await;
        if !resp.status().is_success() {
            let status = resp.status();
            let body = test::read_body(resp).await;
            panic!(
                "test_ingest_zstd_msgpack failed with status {}: {:?}",
                status,
                String::from_utf8_lossy(&body)
            );
        }

        // Verify count
        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM logs")
            .fetch_one(&db_pool)
            .await
            .unwrap();
        assert_eq!(row.0, 10);

        db_pool.close().await;
        let _ = std::fs::remove_file(db_path);
        let _ = std::fs::remove_file(format!("{}-wal", db_path));
        let _ = std::fs::remove_file(format!("{}-shm", db_path));
    }

    #[actix_web::test]
    async fn test_ingest_gzip_json() {
        let db_path = "test_ingest_gzip_json.db";
        let _ = std::fs::remove_file(db_path);

        let db_pool = db::init_db(db_path).await.unwrap();
        let state = web::Data::new(AppState {
            db_pool: db_pool.clone(),
            db_path: db_path.to_string(),
            max_db_size: 1_000_000,
            api_key: None,
        });

        let app = test::init_service(App::new().app_data(state.clone()).service(ingest)).await;

        let logs = make_test_logs(3);
        let json_bytes = serde_json::to_vec(&logs).unwrap();
        let gzip_payload = compress_gzip(&json_bytes);

        let req = test::TestRequest::post()
            .uri("/ingest")
            .insert_header(("Content-Type", "application/json"))
            .insert_header(("X-Content-Encoding", "gzip"))
            .set_payload(gzip_payload)
            .to_request();

        let resp = test::call_service(&app, req).await;
        if !resp.status().is_success() {
            let status = resp.status();
            let body = test::read_body(resp).await;
            panic!(
                "test_ingest_gzip_json failed with status {}: {:?}",
                status,
                String::from_utf8_lossy(&body)
            );
        }

        // Verify count
        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM logs")
            .fetch_one(&db_pool)
            .await
            .unwrap();
        assert_eq!(row.0, 3);

        db_pool.close().await;
        let _ = std::fs::remove_file(db_path);
        let _ = std::fs::remove_file(format!("{}-wal", db_path));
        let _ = std::fs::remove_file(format!("{}-shm", db_path));
    }

    #[actix_web::test]
    async fn test_db_size_pruning() {
        let db_path = "test_db_size_pruning.db";
        let _ = std::fs::remove_file(db_path);

        let db_pool = db::init_db(db_path).await.unwrap();

        let state = web::Data::new(AppState {
            db_pool: db_pool.clone(),
            db_path: db_path.to_string(),
            max_db_size: 100_000, // 100 KB
            api_key: None,
        });

        let app = test::init_service(App::new().app_data(state.clone()).service(ingest)).await;

        // Ingest large batch to exceed 100 KB
        let logs = make_test_logs(1500);
        let json_bytes = serde_json::to_vec(&logs).unwrap();

        let req = test::TestRequest::post()
            .uri("/ingest")
            .insert_header(("Content-Type", "application/json"))
            .set_payload(json_bytes)
            .to_request();

        let resp = test::call_service(&app, req).await;
        assert!(resp.status().is_success());

        let path = std::path::Path::new(db_path);
        assert!(path.exists());
        let file_size = std::fs::metadata(path).unwrap().len();
        println!("Pruned database file size: {} bytes", file_size);
        assert!(
            file_size <= 100_000,
            "Database size {} exceeded limit of 100,000",
            file_size
        );

        // Verify that the row count is less than the ingested 1500, meaning pruning occurred.
        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM logs")
            .fetch_one(&db_pool)
            .await
            .unwrap();
        println!("Pruned database row count: {}", row.0);
        assert!(
            row.0 < 1500,
            "Pruning did not happen, row count is still {}",
            row.0
        );

        db_pool.close().await;
        let _ = std::fs::remove_file(db_path);
        let _ = std::fs::remove_file(format!("{}-wal", db_path));
        let _ = std::fs::remove_file(format!("{}-shm", db_path));
    }

    #[actix_web::test]
    async fn test_api_key_auth() {
        let db_path = "test_api_key_auth.db";
        let _ = std::fs::remove_file(db_path);

        let db_pool = db::init_db(db_path).await.unwrap();
        let state = web::Data::new(AppState {
            db_pool: db_pool.clone(),
            db_path: db_path.to_string(),
            max_db_size: 1_000_000,
            api_key: Some("secret-token-42".to_string()),
        });

        let app = test::init_service(App::new().app_data(state.clone()).service(ingest)).await;

        let logs = make_test_logs(1);
        let payload = serde_json::to_vec(&logs).unwrap();

        // 1. Send request without Authorization header -> should fail with 401
        let req = test::TestRequest::post()
            .uri("/ingest")
            .insert_header(("Content-Type", "application/json"))
            .set_payload(payload.clone())
            .to_request();

        let resp = test::call_service(&app, req).await;
        assert_eq!(resp.status().as_u16(), 401);

        // 2. Send request with incorrect Authorization header -> should fail with 401
        let req = test::TestRequest::post()
            .uri("/ingest")
            .insert_header(("Content-Type", "application/json"))
            .insert_header(("Authorization", "Bearer wrong-token"))
            .set_payload(payload.clone())
            .to_request();

        let resp = test::call_service(&app, req).await;
        assert_eq!(resp.status().as_u16(), 401);

        // 3. Send request with correct Authorization header -> should succeed
        let req = test::TestRequest::post()
            .uri("/ingest")
            .insert_header(("Content-Type", "application/json"))
            .insert_header(("Authorization", "Bearer secret-token-42"))
            .set_payload(payload)
            .to_request();

        let resp = test::call_service(&app, req).await;
        assert!(resp.status().is_success());

        db_pool.close().await;
        let _ = std::fs::remove_file(db_path);
        let _ = std::fs::remove_file(format!("{}-wal", db_path));
        let _ = std::fs::remove_file(format!("{}-shm", db_path));
    }
}