hagfish/src/main.rs

// src/main.rs
mod api;
mod cli;
mod db;
mod ingest;
mod types;

use std::path::Path;
use std::sync::Arc;
use std::time::Instant;

use axum::serve;
use clap::Parser;
use tokio::net::TcpListener;
use tracing_subscriber::EnvFilter;

struct LoggerGuard {
    _guard: Option<tracing_appender::non_blocking::WorkerGuard>,
}

fn init_logging(log_file_path: Option<&str>) -> LoggerGuard {
    let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
    let is_json = std::env::var("HAGFISH_LOG_FORMAT").as_deref() == Ok("json");

    match log_file_path {
        Some(path) => {
            let file = std::fs::File::create(path).expect("Failed to create log file");
            let (non_blocking, guard) = tracing_appender::non_blocking(file);

            let subscriber = tracing_subscriber::fmt()
                .with_env_filter(filter)
                .with_writer(non_blocking);

            if is_json {
                subscriber.json().init();
            } else {
                subscriber.init();
            }

            LoggerGuard {
                _guard: Some(guard),
            }
        }
        None => {
            let subscriber = tracing_subscriber::fmt().with_env_filter(filter);

            if is_json {
                subscriber.json().init();
            } else {
                subscriber.init();
            }

            LoggerGuard { _guard: None }
        }
    }
}

fn cleanup_stale_parquet_files() {
    let tmp_dir = std::env::temp_dir();
    match std::fs::read_dir(&tmp_dir) {
        Ok(entries) => {
            for entry in entries.filter_map(|e| e.ok()) {
                if let Some(name) = entry.file_name().to_str() {
                    if name.starts_with("hagfish-") && name.ends_with(".parquet") {
                        if let Ok(metadata) = entry.metadata() {
                            if let Ok(modified) = metadata.modified() {
                                let age = std::time::SystemTime::now().duration_since(modified);
                                if age.map(|d| d.as_secs() > 3600).unwrap_or(false) {
                                    let _ = std::fs::remove_file(entry.path());
                                }
                            }
                        }
                    }
                }
            }
        }
        Err(e) => tracing::warn!(error = %e, "Failed to list temp dir"),
    }
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let cli = cli::Cli::parse();

    let config = match &cli.config {
        Some(path) => types::Config::load(path)?,
        None => {
            let default_path = std::path::Path::new("config.json");
            if default_path.exists() {
                types::Config::load(default_path)?
            } else {
                eprintln!("WARN: No config file found, using defaults");
                types::Config::default()
            }
        }
    };

    let _logger_guard = init_logging(config.log_file_path.as_deref());

    tracing::info!(
        duckdb_path = %config.duckdb_path,
        bind_address = %config.bind_address,
        "configuration loaded"
    );

    match cli.command {
        cli::Command::Serve => run_serve(&config).await,
        cli::Command::Ingest { full } => run_ingest(&config, full).await,
        cli::Command::Export { out } => run_export(&config, &out),
    }
}

async fn run_serve(config: &types::Config) -> anyhow::Result<()> {
    cleanup_stale_parquet_files();

    let conn = db::init(&config.duckdb_path)?;

    let state = api::AppState {
        conn: Arc::new(tokio::sync::Mutex::new(conn)),
        config: Arc::new(config.clone()),
    };

    let app = api::build_router(state);

    let listener = TcpListener::bind(&config.bind_address).await?;
    tracing::info!(address = %config.bind_address, "hagfish server started");

    serve(listener, app).await?;

    Ok(())
}

async fn run_ingest(config: &types::Config, full: bool) -> anyhow::Result<()> {
    let run_start = Instant::now();

    tracing::info!(
        mode = if full { "full" } else { "incremental" },
        "starting ingestion run"
    );

    let conn = db::init(&config.duckdb_path)?;
    let client = reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(60))
        .build()?;

    let (lookup_maps, metadata) = ingest::fetch_metadata(&client, &config.data_source_url)
        .await
        .inspect_err(|e| {
            tracing::error!(error = %e, stage = "fetch_metadata", "ingestion failed");
        })?;

    db::update_lookups(&conn, &lookup_maps)?;

    let all_months = ingest::extract_available_months(&metadata);

    let months_to_fetch: Vec<String> = if full {
        all_months
    } else {
        let last = db::get_last_month(&conn)?;
        match last {
            Some(last_month) => all_months
                .into_iter()
                .filter(|m| m.as_str() > last_month.as_str())
                .collect(),
            None => all_months,
        }
    };

    if months_to_fetch.is_empty() {
        let elapsed = run_start.elapsed();
        tracing::info!(
            mode = if full { "full" } else { "incremental" },
            months_requested = 0,
            total_rows = 0,
            batches = 0,
            errors = 0,
            elapsed_ms = elapsed.as_millis() as u64,
            "ingestion complete — no new months"
        );
        return Ok(());
    }

    const BATCH_SIZE: usize = 12;
    let batches = ingest::chunk_months(&months_to_fetch, BATCH_SIZE);

    tracing::info!(
        mode = if full { "full" } else { "incremental" },
        months_requested = months_to_fetch.len(),
        batches = batches.len(),
        "beginning batched ingestion"
    );

    let mut total_rows = 0;
    let mut errors = 0;

    for (i, batch) in batches.iter().enumerate() {
        let batch_start = Instant::now();

        tracing::info!(
            batch = i + 1,
            total_batches = batches.len(),
            months_in_batch = batch.len(),
            months = %batch.join(", "),
            "fetching batch"
        );

        let query = ingest::build_query(batch)?;

        let data = match ingest::fetch_data(&client, &config.data_source_url, &query).await {
            Ok(d) => d,
            Err(e) => {
                tracing::error!(
                    batch = i + 1,
                    error = %e,
                    stage = "fetch_data",
                    "batch failed, continuing to next batch"
                );
                errors += 1;
                continue;
            }
        };

        let mut landings = Vec::with_capacity(data.dataset.value.len());
        let mut parse_errors = 0;

        for j in 0..data.dataset.value.len() {
            match ingest::parse_row(j, &data.dataset, &lookup_maps) {
                Ok(row) => {
                    landings.push(ingest::data_row_to_landing(row));
                }
                Err(e) => {
                    tracing::warn!(
                        batch = i + 1,
                        row_index = j,
                        error = %e,
                        "failed to parse row"
                    );
                    parse_errors += 1;
                    errors += 1;
                }
            }
        }

        let count = match db::upsert_landings(&conn, &landings) {
            Ok(c) => c,
            Err(e) => {
                tracing::error!(
                    batch = i + 1,
                    error = %e,
                    stage = "upsert_landings",
                    "batch upsert failed"
                );
                errors += 1;
                continue;
            }
        };

        total_rows += count;
        let batch_elapsed = batch_start.elapsed();

        tracing::info!(
            batch = i + 1,
            total_batches = batches.len(),
            rows_inserted = count,
            parse_errors = parse_errors,
            elapsed_ms = batch_elapsed.as_millis() as u64,
            "batch complete"
        );

        if i < batches.len() - 1 {
            tokio::time::sleep(std::time::Duration::from_secs(1)).await;
        }
    }

    let elapsed = run_start.elapsed();

    tracing::info!(
        mode = if full { "full" } else { "incremental" },
        months_requested = months_to_fetch.len(),
        batches = batches.len(),
        total_rows = total_rows,
        errors = errors,
        elapsed_ms = elapsed.as_millis() as u64,
        "ingestion run complete"
    );

    Ok(())
}

fn run_export(config: &types::Config, out: &Path) -> anyhow::Result<()> {
    let run_start = Instant::now();

    let conn = db::init(&config.duckdb_path)?;
    let path_str = out
        .to_str()
        .ok_or_else(|| anyhow::anyhow!("Output path contains invalid UTF-8"))?;

    db::export_parquet(&conn, path_str)?;

    let elapsed = run_start.elapsed();
    tracing::info!(
        path = path_str,
        elapsed_ms = elapsed.as_millis() as u64,
        "export complete"
    );

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::Config;

    #[test]
    fn test_config_default_values() {
        let config = Config::default();
        assert_eq!(config.duckdb_path, "hagfish.db");
        assert_eq!(config.bind_address, "127.0.0.1:8090");
        assert!(config.data_source_url.contains("statbank.hagstova.fo"));
    }

    #[test]
    fn test_run_export_writes_valid_parquet() {
        let tmp_dir = tempfile::tempdir().unwrap();
        let db_path = tmp_dir.path().join("test.db");
        let parquet_path = tmp_dir.path().join("export.parquet");

        let conn = db::init(db_path.to_str().unwrap()).unwrap();

        let rows = vec![crate::types::Landing {
            month: "2024M01".to_string(),
            species_code: "COD".to_string(),
            species_label: "Toskur".to_string(),
            gear_code: "TOTAL".to_string(),
            zone_code: "TOTAL".to_string(),
            processing_code: "TOTAL".to_string(),
            preservation_code: "TOTAL".to_string(),
            shipsize_code: "TOTAL".to_string(),
            measure_code: "MASS".to_string(),
            value: Some(1000.0),
        }];

        db::upsert_landings(&conn, &rows).unwrap();
        drop(conn);

        let config = Config {
            duckdb_path: db_path.to_str().unwrap().to_string(),
            ..Default::default()
        };

        let result = run_export(&config, &parquet_path);
        assert!(result.is_ok());

        let metadata = std::fs::metadata(&parquet_path).unwrap();
        assert!(metadata.len() > 0);

        let file_bytes = std::fs::read(&parquet_path).unwrap();
        assert_eq!(&file_bytes[..4], b"PAR1");
    }
}