kcode-dev-tools 0.2.0

use std::collections::HashMap;
use std::sync::{Arc, Condvar, Mutex};
use std::time::Instant;

use crate::backend::OpenedSource;
use crate::{ManagedSourceKind, SESSION_LEASE, ToolError};

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
struct HandleKey {
    session_id: String,
    kind: ManagedSourceKind,
    name: String,
}

struct OpenHandle {
    source: Arc<Mutex<OpenedSource>>,
    active_operations: usize,
    closing: bool,
    last_used: Instant,
}

#[derive(Default)]
struct Registry {
    sources: HashMap<HandleKey, OpenHandle>,
}

pub(crate) struct RegistryState {
    entries: Mutex<Registry>,
    changed: Condvar,
}

pub(crate) struct ActiveOperation {
    registry: Arc<RegistryState>,
    key: HandleKey,
}

impl Default for RegistryState {
    fn default() -> Self {
        Self {
            entries: Mutex::new(Registry::default()),
            changed: Condvar::new(),
        }
    }
}

impl RegistryState {
    pub(crate) fn contains(
        &self,
        session_id: &str,
        kind: ManagedSourceKind,
        name: &str,
    ) -> Result<bool, ToolError> {
        let registry = self.lock()?;
        Ok(registry
            .sources
            .contains_key(&HandleKey::new(session_id, kind, name)))
    }

    pub(crate) fn insert(
        &self,
        session_id: &str,
        kind: ManagedSourceKind,
        name: &str,
        source: OpenedSource,
    ) -> Result<(), ToolError> {
        let key = HandleKey::new(session_id, kind, name);
        let mut registry = self.lock()?;
        if registry.sources.contains_key(&key) {
            return Err(ToolError::conflict(
                kind.code("already_open"),
                format!("{} {name:?} is already open in this session.", kind.label()),
            ));
        }
        registry.sources.insert(key, OpenHandle::new(source));
        Ok(())
    }

    pub(crate) fn begin(
        self: &Arc<Self>,
        session_id: &str,
        kind: ManagedSourceKind,
        name: &str,
    ) -> Result<(Arc<Mutex<OpenedSource>>, ActiveOperation), ToolError> {
        let key = HandleKey::new(session_id, kind, name);
        let mut registry = self.lock()?;
        let Some(entry) = registry.sources.get_mut(&key) else {
            return Err(ToolError::conflict(
                kind.code("not_open"),
                format!(
                    "{} {name:?} is not open in this session. Call {} first.",
                    kind.label(),
                    kind.open_tool()
                ),
            ));
        };
        if entry.closing {
            return Err(ToolError::conflict(
                kind.code("session_ending"),
                format!("The session is releasing its {} snapshot.", kind.label()),
            ));
        }
        entry.active_operations += 1;
        entry.last_used = Instant::now();
        Ok((
            Arc::clone(&entry.source),
            ActiveOperation {
                registry: Arc::clone(self),
                key,
            },
        ))
    }

    pub(crate) fn remove_expired(&self) -> Result<(), ToolError> {
        let mut registry = self.lock()?;
        let now = Instant::now();
        registry.sources.retain(|_, entry| {
            entry.active_operations > 0 || now.duration_since(entry.last_used) < SESSION_LEASE
        });
        Ok(())
    }

    pub(crate) fn release(&self, session_id: &str) -> Result<usize, ToolError> {
        let mut registry = self.lock()?;
        let keys = registry
            .sources
            .keys()
            .filter(|key| key.session_id == session_id)
            .cloned()
            .collect::<Vec<_>>();
        for key in &keys {
            if let Some(entry) = registry.sources.get_mut(key) {
                entry.closing = true;
            }
        }
        while keys.iter().any(|key| {
            registry
                .sources
                .get(key)
                .is_some_and(|entry| entry.active_operations > 0)
        }) {
            registry = self
                .changed
                .wait(registry)
                .map_err(|_| ToolError::internal("The source registry is unavailable."))?;
        }
        for key in &keys {
            registry.sources.remove(key);
        }
        Ok(keys.len())
    }

    fn lock(&self) -> Result<std::sync::MutexGuard<'_, Registry>, ToolError> {
        self.entries
            .lock()
            .map_err(|_| ToolError::internal("The source registry is unavailable."))
    }
}

impl Drop for ActiveOperation {
    fn drop(&mut self) {
        let Ok(mut registry) = self.registry.entries.lock() else {
            return;
        };
        let Some(entry) = registry.sources.get_mut(&self.key) else {
            return;
        };
        if entry.active_operations == 0 {
            return;
        }
        entry.active_operations -= 1;
        self.registry.changed.notify_all();
    }
}

impl HandleKey {
    fn new(session_id: &str, kind: ManagedSourceKind, name: &str) -> Self {
        Self {
            session_id: session_id.to_owned(),
            kind,
            name: name.to_owned(),
        }
    }
}

impl OpenHandle {
    fn new(source: OpenedSource) -> Self {
        Self {
            source: Arc::new(Mutex::new(source)),
            active_operations: 0,
            closing: false,
            last_used: Instant::now(),
        }
    }
}