use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, sync::Arc};
use tokio::sync::{RwLock, broadcast};
use tokio_util::sync::CancellationToken;
use uuid::Uuid;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RunEvent {
    pub run_id: Uuid,
    pub event_type: String,
    pub stage: String,
    pub message: String,
    pub progress_current: i64,
    pub progress_total: Option<i64>,
    pub at: DateTime<Utc>,
}

#[derive(Clone)]
pub struct RunControl {
    tokens: Arc<RwLock<HashMap<Uuid, CancellationToken>>>,
    events: broadcast::Sender<RunEvent>,
}

impl Default for RunControl {
    fn default() -> Self {
        Self::new(1_024)
    }
}

impl RunControl {
    pub fn new(event_capacity: usize) -> Self {
        let (events, _) = broadcast::channel(event_capacity.max(16));
        Self {
            tokens: Arc::new(RwLock::new(HashMap::new())),
            events,
        }
    }

    pub async fn token(&self, run_id: Uuid) -> CancellationToken {
        let mut tokens = self.tokens.write().await;
        tokens.entry(run_id).or_default().clone()
    }

    pub async fn cancel(&self, run_id: Uuid) -> bool {
        let token = self.tokens.read().await.get(&run_id).cloned();
        if let Some(token) = token {
            token.cancel();
            true
        } else {
            false
        }
    }

    pub async fn forget(&self, run_id: Uuid) {
        self.tokens.write().await.remove(&run_id);
    }

    pub fn publish(&self, event: RunEvent) {
        let _ = self.events.send(event);
    }

    pub fn subscribe(&self) -> broadcast::Receiver<RunEvent> {
        self.events.subscribe()
    }
}

impl RunEvent {
    pub fn status(run_id: Uuid, stage: impl Into<String>, message: impl Into<String>) -> Self {
        Self {
            run_id,
            event_type: "status".into(),
            stage: stage.into(),
            message: message.into(),
            progress_current: 0,
            progress_total: None,
            at: Utc::now(),
        }
    }
}