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

use anyhow::{Result, anyhow};
use api::{api_complete_json, api_loading_json, api_queue_json};
use axum::{Router, routing::get};
use log::{debug, error, info};
use serde_json::{Map, Value};
use std::{
    collections::{HashMap, HashSet},
    env::args,
    path::Path,
    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>>,

    config: Value,

    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_config().await?;
    state.load().await?;

    let bind_addr = state
        .config
        .get("bind_addr")
        .and_then(Value::as_str)
        .unwrap_or("127.0.0.1");
    let bind_port = state
        .config
        .get("bind_port")
        .and_then(Value::as_u64)
        .unwrap_or(44794) as u16;

    let listener = TcpListener::bind((bind_addr, bind_port)).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)));

    axum::serve(listener, router).await?;
    Ok(())
}

impl Default for State {
    fn default() -> Self {
        Self {
            config: Default::default(),
            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_config(&mut self) -> Result<()> {
        let path = args()
            .nth(1)
            .ok_or(anyhow!("first argument is config path"))?;

        self.config = serde_yml::from_str::<serde_json::Value>(&read_to_string(path).await?)?;
        Ok(())
    }
    pub async fn load(&mut self) -> Result<()> {
        debug!("loading state");
        let t = Instant::now();
        if AsRef::<Path>::as_ref("metadata.json").exists()
            && AsRef::<Path>::as_ref("queue.json").exists()
            && AsRef::<Path>::as_ref("complete.json").exists()
        {
            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());
        } else {
            error!("some state files are missing, skipping load")
        }
        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(())
    }
}