use std::sync::Arc;

use dashmap::DashMap;
use tera::Tera;
use tokio::sync::broadcast;
use uuid::Uuid;

use common_contracts::observability::TaskProgress;

use crate::{config::AppConfig, templates::load_tera};

#[derive(Clone)]
pub struct AppState {
    pub tasks: Arc<DashMap<Uuid, TaskProgress>>,
    pub streams: Arc<StreamManager>,
    pub config: Arc<AppConfig>,
    #[allow(dead_code)]
    pub tera: Arc<Tera>,
    pub nats: async_nats::Client,
}

pub struct StreamManager {
    // Key: RequestId (Uuid)
    // Value: Broadcast Sender. We send JSON strings.
    pub channels: DashMap<Uuid, broadcast::Sender<String>>, 
}

impl StreamManager {
    pub fn new() -> Self {
        Self {
            channels: DashMap::new(),
        }
    }
    
    pub fn get_or_create_sender(&self, request_id: Uuid) -> broadcast::Sender<String> {
        self.channels.entry(request_id).or_insert_with(|| {
             // Buffer size 100 messages. If lag occurs, receiver might miss messages (RecvError::Lagged).
             // Since we use this for real-time text streaming, missing a chunk is bad but acceptable for "pseudo" visual if we can recover,
             // but here we don't want to miss content. 
             // Frontend SSE usually just reconnects. 
             // We should probably use a larger buffer to be safe, e.g. 1000.
             let (tx, _rx) = broadcast::channel(1000);
             tx
        }).value().clone()
    }

    #[allow(dead_code)]
    pub fn remove_channel(&self, request_id: &Uuid) {
        self.channels.remove(request_id);
    }
}

impl AppState {
    pub fn new(config: AppConfig, nats: async_nats::Client) -> Self {
        Self {
            tasks: Arc::new(DashMap::new()),
            streams: Arc::new(StreamManager::new()),
            config: Arc::new(config),
            tera: Arc::new(load_tera()),
            nats,
        }
    }
}