use dashmap::DashMap;
use tokio::sync::{mpsc, oneshot};
use uuid::Uuid;
use serde::Serialize;
use helios_common::protocol::{ClientMessage, ServerMessage};

/// Represents one connected remote client
#[derive(Debug, Clone)]
pub struct Session {
    pub id: Uuid,
    pub label: Option<String>,
    /// Channel to send commands to the WS handler for this session
    pub cmd_tx: mpsc::Sender<ServerMessage>,
}

/// Serializable view of a session for the REST API
#[derive(Debug, Serialize)]
pub struct SessionInfo {
    pub id: Uuid,
    pub label: Option<String>,
}

impl From<&Session> for SessionInfo {
    fn from(s: &Session) -> Self {
        SessionInfo {
            id: s.id,
            label: s.label.clone(),
        }
    }
}

pub struct SessionStore {
    /// Active sessions by ID
    sessions: DashMap<Uuid, Session>,
    /// Pending request callbacks by request_id
    pending: DashMap<Uuid, oneshot::Sender<ClientMessage>>,
}

impl SessionStore {
    pub fn new() -> Self {
        Self {
            sessions: DashMap::new(),
            pending: DashMap::new(),
        }
    }

    pub fn insert(&self, session: Session) {
        self.sessions.insert(session.id, session);
    }

    pub fn remove(&self, id: &Uuid) {
        self.sessions.remove(id);
    }

    pub fn get_cmd_tx(&self, id: &Uuid) -> Option<mpsc::Sender<ServerMessage>> {
        self.sessions.get(id).map(|s| s.cmd_tx.clone())
    }

    pub fn set_label(&self, id: &Uuid, label: String) -> bool {
        if let Some(mut s) = self.sessions.get_mut(id) {
            s.label = Some(label);
            true
        } else {
            false
        }
    }

    pub fn list(&self) -> Vec<SessionInfo> {
        self.sessions.iter().map(|e| SessionInfo::from(e.value())).collect()
    }

    /// Register a pending request. Returns the receiver to await the client response.
    pub fn register_pending(&self, request_id: Uuid) -> oneshot::Receiver<ClientMessage> {
        let (tx, rx) = oneshot::channel();
        self.pending.insert(request_id, tx);
        rx
    }

    /// Deliver a client response to the waiting request handler.
    /// Returns true if the request was found and resolved.
    pub fn resolve_pending(&self, request_id: Uuid, msg: ClientMessage) -> bool {
        if let Some((_, tx)) = self.pending.remove(&request_id) {
            let _ = tx.send(msg);
            true
        } else {
            false
        }
    }
}