use std::{collections::HashMap, sync::Arc};

use tokio::sync::Mutex;

use crate::{config::AppConfig, inference_process::InferenceProcess, util::parse_size};

#[derive(Clone, Debug)]
pub struct ResourceManager {
    pub total_ram:  u64,
    pub total_vram: u64,
    pub used_ram:   u64,
    pub used_vram:  u64,
    pub processes:  HashMap<u16, InferenceProcess>,
}

pub type ResourceManagerHandle = Arc<Mutex<ResourceManager>>;

#[derive(Clone, Debug, derive_more::Deref)]
pub struct AppState(pub ResourceManagerHandle);

impl AppState {
    pub fn from_config(config: AppConfig) -> Self {
        let total_ram =
            parse_size(&config.system_resources.ram).expect("Invalid RAM size in config");
        let total_vram =
            parse_size(&config.system_resources.vram).expect("Invalid VRAM size in config");

        let resource_manager = ResourceManager {
            total_ram,
            total_vram,
            used_ram: 0,
            used_vram: 0,
            processes: HashMap::new(),
        };

        Self(Arc::new(Mutex::new(resource_manager)))
    }
}