crates/bex-core/src/registry.rs

use bex_types::plugin_info::PluginInfo;
use bex_types::{Capabilities, Manifest};
use indexmap::IndexMap;
use std::sync::atomic::{AtomicU32, AtomicU64, AtomicU8, Ordering};
use std::sync::Arc;
use wasmtime::component::Component;

/// Health tracking for circuit breaker pattern.
/// States: 0 = closed (healthy), 1 = open (broken), 2 = half-open (testing)
pub struct PluginHealth {
    consecutive_failures: AtomicU32,
    last_failure_at: AtomicU64,
    circuit_state: AtomicU8,
    threshold: u32,
    cooldown_ms: u64,
}

impl PluginHealth {
    pub fn new(threshold: u32, cooldown_ms: u64) -> Self {
        Self {
            consecutive_failures: AtomicU32::new(0),
            last_failure_at: AtomicU64::new(0),
            circuit_state: AtomicU8::new(0),
            threshold,
            cooldown_ms,
        }
    }

    pub fn record_success(&self) {
        self.consecutive_failures.store(0, Ordering::Relaxed);
        self.circuit_state.store(0, Ordering::Relaxed);
    }

    pub fn record_failure(&self) {
        let count = self.consecutive_failures.fetch_add(1, Ordering::Relaxed) + 1;
        self.last_failure_at.store(now_ms(), Ordering::Relaxed);
        if count >= self.threshold {
            self.circuit_state.store(1, Ordering::Relaxed);
        }
    }

    pub fn is_available(&self) -> bool {
        match self.circuit_state.load(Ordering::Relaxed) {
            0 => true,
            1 => {
                if now_ms() - self.last_failure_at.load(Ordering::Relaxed) > self.cooldown_ms {
                    self.circuit_state.store(2, Ordering::Relaxed);
                    true
                } else {
                    false
                }
            }
            2 => true,
            _ => false,
        }
    }
}

fn now_ms() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as u64
}

pub struct PluginRecord {
    pub id: String,
    pub manifest: Arc<Manifest>,
    pub capabilities: Capabilities,
    pub enabled: bool,
    pub installed_at: u64,
    pub component: Arc<Component>,
    pub health: PluginHealth,
}

impl PluginRecord {
    pub fn new(manifest: Manifest, component: Arc<Component>, threshold: u32, cooldown_ms: u64) -> Self {
        let id = manifest.id.clone();
        let capabilities = manifest.capabilities();
        let installed_at = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_millis() as u64;
        Self {
            id,
            manifest: Arc::new(manifest),
            capabilities,
            enabled: true,
            installed_at,
            component,
            health: PluginHealth::new(threshold, cooldown_ms),
        }
    }

    pub fn to_plugin_info(&self) -> PluginInfo {
        PluginInfo {
            id: self.id.clone(),
            name: self.manifest.name.clone(),
            version: self.manifest.version.clone(),
            capabilities: self.capabilities.bits(),
            description: self.manifest.display.description.clone().unwrap_or_default(),
            tags: self.manifest.display.tags.clone(),
            priority: self.manifest.display.priority,
            enabled: self.enabled,
            installed_at: self.installed_at,
        }
    }

    /// Returns sorted-by-priority list of enabled plugins with the given capability
    pub fn matches_capability(&self, cap: Capabilities) -> bool {
        self.enabled && self.capabilities.has(cap)
    }
}

/// Plugin registry with deterministic insertion order via IndexMap.
pub struct PluginRegistry {
    by_id: IndexMap<String, Arc<PluginRecord>>,
}

impl PluginRegistry {
    pub fn new() -> Self {
        Self {
            by_id: IndexMap::new(),
        }
    }

    pub fn insert(&mut self, record: Arc<PluginRecord>) {
        self.by_id.insert(record.id.clone(), record);
    }

    pub fn remove(&mut self, id: &str) -> Option<Arc<PluginRecord>> {
        self.by_id.shift_remove(id)
    }

    pub fn get(&self, id: &str) -> Option<&Arc<PluginRecord>> {
        self.by_id.get(id)
    }

    /// Returns list in deterministic insertion order
    pub fn list(&self) -> Vec<PluginInfo> {
        self.by_id.values().map(|r| r.to_plugin_info()).collect()
    }

    /// Returns plugins sorted by priority (highest first), then by ID for stable tiebreak
    pub fn list_sorted_by_priority(&self) -> Vec<Arc<PluginRecord>> {
        let mut plugins: Vec<_> = self.by_id.values().cloned().collect();
        plugins.sort_by(|a, b| {
            b.manifest.display.priority
                .cmp(&a.manifest.display.priority)
                .then_with(|| a.id.cmp(&b.id))
        });
        plugins
    }

    /// List enabled plugins with the given capability, sorted by priority
    pub fn list_with_cap(&self, cap: Capabilities) -> Vec<Arc<PluginRecord>> {
        let mut plugins: Vec<_> = self
            .by_id
            .values()
            .filter(|r| r.matches_capability(cap))
            .cloned()
            .collect();
        plugins.sort_by(|a, b| {
            b.manifest.display.priority
                .cmp(&a.manifest.display.priority)
                .then_with(|| a.id.cmp(&b.id))
        });
        plugins
    }

    /// Get mutable reference — needed for enable/disable
    pub fn get_arc_mut(&mut self, id: &str) -> Option<&mut Arc<PluginRecord>> {
        self.by_id.get_mut(id)
    }
}