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

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use tokio::sync::{Notify, RwLock};
use uuid::Uuid;

use super::state::SegmentState;

/// Per-segment runtime as defined in design doc §4.2.6.
///
/// Held in memory only; reset on restart.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SegmentRuntime {
    pub segment_id: Uuid,
    pub state: SegmentState,
    pub auto_enabled: bool,
    pub current_step_no: Option<i32>,
    pub step_started_at: Option<DateTime<Utc>>,
    pub last_completed_at: Option<DateTime<Utc>>,
    pub blocked_reason: Option<String>,
    pub fault_message: Option<String>,
    pub manual_ack_required: bool,
    pub comm_locked: bool,
    pub rem_local: bool,
    pub held_resources: Vec<String>,
}

impl SegmentRuntime {
    pub fn new(segment_id: Uuid) -> Self {
        Self {
            segment_id,
            state: SegmentState::Idle,
            auto_enabled: false,
            current_step_no: None,
            step_started_at: None,
            last_completed_at: None,
            blocked_reason: None,
            fault_message: None,
            manual_ack_required: false,
            comm_locked: false,
            rem_local: false,
            held_resources: Vec::new(),
        }
    }
}

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

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

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

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

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

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

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

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

    pub async fn notify_segment(&self, segment_id: Uuid) {
        if let Some(notify) = self.notifiers.read().await.get(&segment_id) {
            notify.notify_one();
        }
    }
}