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

use tokio::sync::{Notify, RwLock};
use uuid::Uuid;

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum UnitRuntimeState {
    Stopped,
    Running,
    DistributorRunning,
    FaultLocked,
    CommLocked,
}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct UnitRuntime {
    pub unit_id: Uuid,
    pub state: UnitRuntimeState,
    pub auto_enabled: bool,
    pub accumulated_run_sec: i64,
    /// Snapshot updated only on state transitions; used for display to avoid mid-tick jitter.
    pub display_acc_sec: i64,
    pub fault_locked: bool,
    pub flt_active: bool,
    pub comm_locked: bool,
    pub manual_ack_required: bool,
}

impl UnitRuntime {
    pub fn new(unit_id: Uuid) -> Self {
        Self {
            unit_id,
            state: UnitRuntimeState::Stopped,
            auto_enabled: false,
            accumulated_run_sec: 0,
            display_acc_sec: 0,
            fault_locked: false,
            flt_active: false,
            comm_locked: false,
            manual_ack_required: false,
        }
    }
}

#[derive(Clone, Default)]
pub struct ControlRuntimeStore {
    inner: Arc<RwLock<HashMap<Uuid, UnitRuntime>>>,
    notifiers: Arc<RwLock<HashMap<Uuid, Arc<Notify>>>>,
}

impl ControlRuntimeStore {
    pub fn new() -> Self {
        Self::default()
    }

    pub async fn get(&self, unit_id: Uuid) -> Option<UnitRuntime> {
        self.inner.read().await.get(&unit_id).cloned()
    }

    pub async fn get_or_init(&self, unit_id: Uuid) -> UnitRuntime {
        if let Some(runtime) = self.get(unit_id).await {
            return runtime;
        }

        let runtime = UnitRuntime::new(unit_id);
        self.inner.write().await.insert(unit_id, runtime.clone());
        runtime
    }

    pub async fn upsert(&self, runtime: UnitRuntime) {
        self.inner.write().await.insert(runtime.unit_id, runtime);
    }

    pub async fn get_or_create_notify(&self, unit_id: Uuid) -> Arc<Notify> {
        self.notifiers
            .write()
            .await
            .entry(unit_id)
            .or_insert_with(|| Arc::new(Notify::new()))
            .clone()
    }

    pub async fn get_all(&self) -> HashMap<Uuid, UnitRuntime> {
        self.inner.read().await.clone()
    }

    /// Wake the engine task for a unit (e.g., when auto_enabled or fault_locked changes).
    pub async fn notify_unit(&self, unit_id: Uuid) {
        if let Some(n) = self.notifiers.read().await.get(&unit_id) {
            n.notify_one();
        }
    }
}