poop-graph/src/main.rs

use std::{path::PathBuf, sync::Arc, time::Duration};

use actix_web::{
    get, http::header::ContentType, middleware::Logger, web, App, HttpResponse, HttpServer,
    Responder,
};
use lazy_static::lazy_static;
use tokio::sync::RwLock;

mod extract;
mod fetch;
mod serialize;

use extract::DataSet;
use fetch::PdfFetcher;
use serialize::{Csv, DataSerializer, Json};

lazy_static! {
    static ref UPDATE_INTERVAL: Duration = Duration::from_secs(3600);
}

struct AppState {
    dataset: RwLock<Arc<DataSet>>,
}

#[derive(thiserror::Error, Debug)]
enum Error {
    #[error("Failed to fetch PDF")]
    Fetch(#[from] fetch::Error),
    #[error("Failed to extract data from PDF")]
    Extract(#[from] extract::Error),
}

async fn load_data(fetcher: &mut PdfFetcher) -> Result<DataSet, Error> {
    Ok(DataSet::extract(&fetcher.fetch().await?)?)
}

async fn try_update(state: &AppState, fetcher: &mut PdfFetcher) -> Result<(), Error> {
    *state.dataset.write().await = Arc::new(load_data(fetcher).await?);
    Ok(())
}

async fn start_updater() -> Result<web::Data<AppState>, Error> {
    let cached_pdf_path = PathBuf::from(std::env::var("CACHED_PDF_PATH").unwrap_or_else(|_| String::from("data.pdf")));
    let mut fetcher = PdfFetcher::new(cached_pdf_path)?;
    let state = web::Data::new(AppState {
        dataset: RwLock::new(Arc::new(load_data(&mut fetcher).await?)),
    });

    let state_copy = state.clone();
    std::thread::spawn(move || {
        actix_web::rt::System::new().block_on(async {
            loop {
                actix_web::rt::time::sleep(*UPDATE_INTERVAL).await;
                if let Err(e) = try_update(&state_copy, &mut fetcher).await {
                    eprintln!("Error updating data: {:#?}", e);
                }
            }
        });
    });

    Ok(state)
}

#[actix_web::main]
async fn main() -> std::io::Result<()> {
    simple_logger::init_with_level(log::Level::Info).unwrap();

    let static_dir = PathBuf::from(std::env::var("STATIC_DIR").unwrap_or_else(|_| String::from("static")));

    let state = start_updater().await.expect("Failed to initialize state");

    HttpServer::new(move || {
        App::new()
            .app_data(state.clone())
            .wrap(Logger::default())
            .service(csv)
            .service(json)
            .service(actix_files::Files::new("/", static_dir.clone()).index_file("index.html"))
    })
    .bind("127.0.0.1:8080")?
    .run()
    .await
}

#[get("/data.csv")]
async fn csv(data: web::Data<AppState>) -> impl Responder {
    let dataset = { data.dataset.read().await.clone() };

    let rows = tokio_stream::iter(
        DataSerializer::new(dataset, Csv).map(|item| item.map(bytes::Bytes::from)),
    );
    HttpResponse::Ok()
        .content_type("text/csv; charset=utf-8")
        .body(actix_web::body::BodyStream::new(rows))
}

#[get("/data.json")]
async fn json(data: web::Data<AppState>) -> impl Responder {
    let dataset = { data.dataset.read().await.clone() };

    let rows = tokio_stream::iter(
        DataSerializer::new(dataset, Json).map(|item| item.map(bytes::Bytes::from)),
    );
    HttpResponse::Ok()
        .insert_header(ContentType::json())
        .body(actix_web::body::BodyStream::new(rows))
}