pub mod api;
pub mod helper;
pub mod webui;
pub mod webui_ws;
pub mod worker_ws;

use anyhow::Result;
use api::{api_complete_json, api_loading_json, api_queue_json};
use axum::{Router, routing::get};
use log::{debug, info};
use serde_json::{Map, Value};
use std::{
    collections::{HashMap, HashSet},
    sync::Arc,
    time::Instant,
};
use tokio::{
    fs::{File, read_to_string, rename},
    io::AsyncWriteExt,
    net::TcpListener,
    sync::{
        RwLock, broadcast,
        mpsc::{self},
    },
};
use webui::{webui, webui_script, webui_style};
use webui_ws::{WebuiEvent, webui_websocket};
use worker_ws::{WorkerID, WorkerResponse, worker_websocket};

pub struct Worker {
    accept: usize,
    name: String,
    task_kinds: Vec<String>,
    assigned_tasks: HashSet<String>,
    send: mpsc::Sender<WorkerResponse>,
}

pub struct State {
    worker_id_counter: WorkerID,
    workers: HashMap<WorkerID, Worker>,
    webui_broadcast: broadcast::Sender<Arc<WebuiEvent>>,

    metadata: HashMap<String, Map<String, Value>>,
    queue: HashSet<String>,
    loading: HashSet<String>,
    complete: HashSet<String>,
}

#[tokio::main]
async fn main() -> Result<()> {
    env_logger::init_from_env("LOG");
    let mut state = State::default();
    state.load().await?;
    let router = Router::new()
        .route("/", get(webui))
        .route("/style.css", get(webui_style))
        .route("/webui_live.js", get(webui_script))
        .route("/webui_ws", get(webui_websocket))
        .route("/worker_ws", get(worker_websocket))
        .route("/api/queue.json", get(api_queue_json))
        .route("/api/complete.json", get(api_complete_json))
        .route("/api/loading.json", get(api_loading_json))
        .with_state(Arc::new(RwLock::new(state)));
    let listener = TcpListener::bind("127.0.0.1:8080").await?;
    axum::serve(listener, router).await?;
    Ok(())
}

impl Default for State {
    fn default() -> Self {
        Self {
            worker_id_counter: Default::default(),
            workers: Default::default(),
            webui_broadcast: broadcast::channel(1024).0,
            metadata: Default::default(),
            queue: Default::default(),
            loading: Default::default(),
            complete: Default::default(),
        }
    }
}

impl State {
    pub async fn load(&mut self) -> Result<()> {
        debug!("loading state");
        let t = Instant::now();
        self.metadata = serde_json::from_str(&read_to_string("metadata.json").await?)?;
        self.queue = serde_json::from_str(&read_to_string("queue.json").await?)?;
        self.complete = serde_json::from_str(&read_to_string("complete.json").await?)?;
        info!("state loaded (took {:?})", t.elapsed());
        Ok(())
    }
    pub async fn save(&mut self) -> Result<()> {
        let t = Instant::now();
        File::create("metadata.json~")
            .await?
            .write_all(&serde_json::to_vec(&self.metadata)?)
            .await?;
        File::create("queue.json~")
            .await?
            .write_all(&serde_json::to_vec(&self.queue)?)
            .await?;
        File::create("complete.json~")
            .await?
            .write_all(&serde_json::to_vec(&self.complete)?)
            .await?;

        rename("metadata.json~", "metadata.json").await?;
        rename("queue.json~", "queue.json").await?;
        rename("complete.json~", "complete.json").await?;
        info!("state saved (took {:?})", t.elapsed());
        Ok(())
    }
}