//! State containers and task scheduling for engine-side management.
//!
//! These types wrap synchronization primitives so that engine code
//! never directly touches `std::sync::RwLock`, `tokio::sync::Mutex`, etc.
//! Guardian E203 enforces this boundary.
//!
//! Task scheduling: `spawn_loop` / `shutdown_all_loops` provide task lifecycle.
//! `LoopHandler` trait + `LoopRegistry` provide declarative loop management.
//! direct `tokio::spawn` usage in engine code.

use std::collections::{HashMap, HashSet};
use std::future::Future;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use crate::critical::{Transaction, TxScope};
use crate::sealed::Sealed;

// ---------------------------------------------------------------------------
// AtomicFlag — simple boolean signal (shutdown, connected, etc.)
// ---------------------------------------------------------------------------

/// A thread-safe boolean flag. Wraps `Arc<AtomicBool>` with semantic methods.
#[derive(Clone)]
pub(crate) struct AtomicFlag(Arc<AtomicBool>);

impl AtomicFlag {
    /// Create a new flag, initially `false`.
    pub fn new() -> Self {
        Self(Arc::new(AtomicBool::new(false)))
    }

    /// Create from an existing `Arc<AtomicBool>` (e.g. from hat framework types).
    pub fn from_raw(raw: Arc<AtomicBool>) -> Self {
        Self(raw)
    }

    /// Set the flag to `true`.
    pub fn set(&self) {
        self.0.store(true, Ordering::SeqCst);
    }

    /// Read the flag value.
    pub fn get(&self) -> bool {
        self.0.load(Ordering::SeqCst)
    }

    /// Reset the flag to `false`.
    pub fn clear(&self) {
        self.0.store(false, Ordering::SeqCst);
    }

    /// Check and clear — returns `true` if the flag was set, then resets to `false`.
    pub fn take(&self) -> bool {
        self.0.swap(false, Ordering::SeqCst)
    }
}

impl Default for AtomicFlag {
    fn default() -> Self {
        Self::new()
    }
}

// ---------------------------------------------------------------------------
// ConnectionCache<V> — keyed resource cache with short-lived locks
// ---------------------------------------------------------------------------

/// A keyed cache for cloneable resources (DB connections, etc.).
///
/// Uses `std::sync::RwLock` internally — lock durations are trivially short
/// (HashMap lookup/insert only), safe for use in async context.
#[derive(Clone)]
pub(crate) struct ConnectionCache<V: Clone> {
    inner: Arc<std::sync::RwLock<HashMap<String, V>>>,
}

impl<V: Clone> ConnectionCache<V> {
    /// Create an empty cache.
    pub fn new() -> Self {
        Self {
            inner: Arc::new(std::sync::RwLock::new(HashMap::new())),
        }
    }

    /// Try to get a cached value (read lock, fast path).
    pub fn get(&self, key: &str) -> Option<V> {
        self.inner.read().unwrap().get(key).cloned()
    }

    /// Insert a value (write lock).
    pub fn insert(&self, key: &str, value: V) {
        self.inner.write().unwrap().insert(key.to_owned(), value);
    }

    /// Remove a value (write lock).
    pub fn remove(&self, key: &str) -> Option<V> {
        self.inner.write().unwrap().remove(key)
    }

    /// Double-checked insert: returns existing value if another task inserted
    /// between the caller's read-miss and this write. Otherwise inserts and
    /// returns the new value.
    pub fn get_or_insert(&self, key: &str, value: V) -> V {
        let mut map = self.inner.write().unwrap();
        if let Some(existing) = map.get(key) {
            return existing.clone();
        }
        map.insert(key.to_owned(), value.clone());
        value
    }
}

impl<V: Clone> Default for ConnectionCache<V> {
    fn default() -> Self {
        Self::new()
    }
}

// ---------------------------------------------------------------------------
// AsyncCell<T> — tokio RwLock wrapper
// ---------------------------------------------------------------------------

/// Async read-write cell. Wraps `tokio::sync::RwLock<T>`.
pub(crate) struct AsyncCell<T>(tokio::sync::RwLock<T>);

impl<T> AsyncCell<T> {
    /// Create a new cell with the given value.
    pub fn new(value: T) -> Self {
        Self(tokio::sync::RwLock::new(value))
    }

    /// Acquire a read guard.
    pub async fn read(&self) -> tokio::sync::RwLockReadGuard<'_, T> {
        self.0.read().await
    }

    /// Acquire a write guard.
    pub async fn write(&self) -> tokio::sync::RwLockWriteGuard<'_, T> {
        self.0.write().await
    }
}

// ---------------------------------------------------------------------------
// AsyncLock<T> — tokio Mutex wrapper
// ---------------------------------------------------------------------------

/// Async mutex. Wraps `tokio::sync::Mutex<T>`.
pub(crate) struct AsyncLock<T>(tokio::sync::Mutex<T>);

impl<T> AsyncLock<T> {
    /// Create a new lock with the given value.
    pub fn new(value: T) -> Self {
        Self(tokio::sync::Mutex::new(value))
    }

    /// Acquire the lock.
    pub async fn lock(&self) -> tokio::sync::MutexGuard<'_, T> {
        self.0.lock().await
    }
}

// ---------------------------------------------------------------------------
// channel — re-export tokio mpsc types
// ---------------------------------------------------------------------------

/// Async channel types (re-exported from tokio).
pub mod channel {
    pub use tokio::sync::mpsc::{
        unbounded_channel, UnboundedReceiver, UnboundedSender,
    };
}

// ---------------------------------------------------------------------------
// CancelToken — cancellation signal for spawn_loop tasks
// ---------------------------------------------------------------------------

/// A cancellation token with both synchronous polling and async waiting.
///
/// Created internally by `spawn_loop` and passed to the task closure.
/// Engine code uses `get()` for polling in loops or `cancelled().await`
/// for async waiting in `select!` branches.
#[derive(Clone)]
pub struct CancelToken {
    flag: Arc<AtomicBool>,
    notify: Arc<tokio::sync::Notify>,
}

impl CancelToken {
    /// Create a new token (not yet cancelled).
    fn new() -> Self {
        Self {
            flag: Arc::new(AtomicBool::new(false)),
            notify: Arc::new(tokio::sync::Notify::new()),
        }
    }

    /// Check if cancellation has been requested (synchronous, for loop polling).
    pub fn get(&self) -> bool {
        self.flag.load(Ordering::SeqCst)
    }

    /// Wait until cancellation is requested (async, for `select!` branches).
    ///
    /// Returns immediately if already cancelled.
    pub async fn cancelled(&self) {
        // Fast path: already cancelled
        if self.flag.load(Ordering::SeqCst) {
            return;
        }
        // Register for notification, then double-check flag
        let notified = self.notify.notified();
        if self.flag.load(Ordering::SeqCst) {
            return;
        }
        notified.await;
    }

    /// Trigger cancellation. Hat-internal only.
    pub(crate) fn trigger(&self) {
        self.flag.store(true, Ordering::SeqCst);
        self.notify.notify_waiters();
    }
}

// ---------------------------------------------------------------------------
// LoopHandle — handle returned by spawn_loop
// ---------------------------------------------------------------------------

/// Handle to a registered loop task. Provides lifecycle control.
pub struct LoopHandle {
    join: tokio::task::JoinHandle<()>,
    cancel: CancelToken,
}

impl LoopHandle {
    /// Check if the task has finished (exited or panicked).
    pub fn is_finished(&self) -> bool {
        self.join.is_finished()
    }

    /// Request graceful cancellation (sets the cancel token).
    /// The task should check `cancel.get()` or `cancel.cancelled().await`.
    pub fn cancel(&self) {
        self.cancel.trigger();
    }

    /// Force-abort the task (immediate, may leave resources in inconsistent state).
    pub fn abort(&self) {
        self.join.abort();
    }

    /// Wait for the task to complete. Consumes the handle.
    pub async fn join(self) {
        let _ = self.join.await;
    }
}

// ---------------------------------------------------------------------------
// Loop registry — internal tracking for shutdown_all_loops
// ---------------------------------------------------------------------------

struct RegistryEntry {
    #[allow(dead_code)]
    name: String,
    cancel: CancelToken,
}

static LOOP_REGISTRY: std::sync::OnceLock<std::sync::Mutex<Vec<RegistryEntry>>> =
    std::sync::OnceLock::new();

fn get_loop_registry() -> &'static std::sync::Mutex<Vec<RegistryEntry>> {
    LOOP_REGISTRY.get_or_init(|| std::sync::Mutex::new(Vec::new()))
}

// ---------------------------------------------------------------------------
// spawn_loop — registered task with injected CancelToken
// ---------------------------------------------------------------------------

/// Spawn a long-running loop task with an injected `CancelToken`.
///
/// The task is registered in the global loop registry so that
/// `shutdown_all_loops()` can trigger all cancel tokens at once.
///
/// Returns a `LoopHandle` for lifecycle control (is_finished, cancel, abort, join).
///
/// # Example
/// ```ignore
/// let handle = mad_hatter::spawn_loop("my-worker", |cancel| async move {
///     loop {
///         if cancel.get() { break; }
///         // do work...
///         tokio::time::sleep(Duration::from_secs(1)).await;
///     }
/// });
/// ```
pub(crate) fn spawn_loop<F, Fut>(name: impl Into<String>, f: F) -> LoopHandle
where
    F: FnOnce(CancelToken) -> Fut,
    Fut: Future<Output = ()> + Send + 'static,
{
    let cancel = CancelToken::new();
    let cancel_for_task = cancel.clone();
    let future = f(cancel_for_task);
    let join = tokio::spawn(future);

    let entry = RegistryEntry {
        name: name.into(),
        cancel: cancel.clone(),
    };
    get_loop_registry().lock().unwrap().push(entry);

    LoopHandle { join, cancel }
}

// ---------------------------------------------------------------------------
// ---------------------------------------------------------------------------
// shutdown_all_loops — trigger all registered cancel tokens
// ---------------------------------------------------------------------------

/// Trigger cancellation for all tasks registered via `spawn_loop`.
///
/// Each task's `CancelToken` is triggered, which:
/// - Makes `cancel.get()` return `true`
/// - Wakes any pending `cancel.cancelled().await`
///
/// Tasks should check their cancel token and exit gracefully.
pub fn shutdown_all_loops() {
    if let Some(registry) = LOOP_REGISTRY.get() {
        let entries = registry.lock().unwrap();
        for entry in entries.iter() {
            entry.cancel.trigger();
        }
    }
}




// ---------------------------------------------------------------------------
// __start_framework — auto-appended by #[mad_hatter::main] macro
// ---------------------------------------------------------------------------

/// Wait for SIGTERM/SIGINT, then trigger graceful shutdown.
///
/// **Do not call directly** — the `#[mad_hatter::main]` macro appends this
/// automatically after the user's function body.
#[doc(hidden)]
pub async fn __start_framework() {
    // Wait for shutdown signal (Ctrl+C / SIGINT)
    tokio::signal::ctrl_c().await.ok();

    // Trigger graceful shutdown of all registered loops
    shutdown_all_loops();

    // Grace period for loops to exit cleanly
    tokio::time::sleep(std::time::Duration::from_secs(3)).await;
}

// ---------------------------------------------------------------------------
// BackoffConfig — exponential backoff parameters
// ---------------------------------------------------------------------------

/// Configuration for exponential backoff in loop handlers.
#[derive(Debug, Clone)]
pub struct BackoffConfig {
    base_secs: u64,
    max_secs: u64,
}

impl BackoffConfig {
    /// Create exponential backoff config with base and maximum seconds.
    pub fn new(base_secs: u64, max_secs: u64) -> Self {
        Self { base_secs, max_secs }
    }

    fn calculate(&self, consecutive_failures: u32) -> Duration {
        let multiplier = 1u64.checked_shl(consecutive_failures).unwrap_or(u64::MAX);
        let secs = self.base_secs.saturating_mul(multiplier).min(self.max_secs);
        Duration::from_secs(secs)
    }
}

impl Default for BackoffConfig {
    fn default() -> Self {
        Self { base_secs: 2, max_secs: 120 }
    }
}

// ---------------------------------------------------------------------------
// LoopAction / RecoveryAction — loop control enums
// ---------------------------------------------------------------------------

/// Action returned by `LoopHandler::tick` to control the loop.
pub enum LoopAction {
    /// Immediately proceed to the next iteration.
    Continue,
    /// Sleep for the given duration before the next iteration (cancel-aware).
    Sleep(Duration),
    /// Exit the loop.
    Exit,
}

/// Action returned by `LoopHandler::on_fatal` / `on_unexpected_panic`.
pub enum RecoveryAction {
    /// Reset backoff counter, continue immediately.
    Continue,
    /// Framework computes exponential backoff sleep.
    Backoff,
    /// Exit the loop.
    Exit,
}

// ---------------------------------------------------------------------------
// LoopHandler — declarative loop trait
// ---------------------------------------------------------------------------

/// Declarative trait for loop tasks. The framework manages:
/// - `TxScope` creation per iteration
/// - `catch_unwind` for panic recovery
/// - Exponential backoff on failures
/// - Cancel-aware sleep
///
/// Engine implements this trait; the framework runs the loop.
#[async_trait]
pub trait LoopHandler: Send + 'static {
    /// The transaction type for this loop's scope.
    type Tx: Transaction;

    /// Core iteration logic. Called once per loop iteration with a fresh TxScope.
    async fn tick(&mut self, scope: &TxScope<Self::Tx>) -> LoopAction;

    /// Called when tick panics with a known fatal token.
    fn on_fatal(&mut self, token: <Self::Tx as Transaction>::Fatal, detail: Sealed<String>) -> RecoveryAction;

    /// Called when tick panics without a known fatal token. Default: backoff.
    fn on_unexpected_panic(&mut self) -> RecoveryAction {
        RecoveryAction::Backoff
    }
}

/// Internal: run a LoopHandler inside a spawn_loop task.
#[doc(hidden)]
pub fn run_handler_loop<H: LoopHandler>(name: String, mut handler: H, db: Option<crate::db::Db>) -> LoopHandle {
    let backoff = BackoffConfig::default();
    spawn_loop(name, |cancel| async move {
        let mut consecutive_failures: u32 = 0;
        loop {
            if cancel.get() { break; }
            let scope = match &db {
                Some(d) => TxScope::<H::Tx>::new_with_db(d.clone()),
                None => TxScope::<H::Tx>::new(),
            };
            let outcome = crate::critical::run_tx_scope(&scope, async {
                handler.tick(&scope).await
            }).await;

            let action = match outcome {
                crate::critical::TxOutcome::Completed(loop_action) => {
                    consecutive_failures = 0;
                    loop_action
                }
                crate::critical::TxOutcome::Aborted(token, detail) => {
                    let recovery = handler.on_fatal(token, detail);
                    match recovery {
                        RecoveryAction::Continue => {
                            consecutive_failures = 0;
                            LoopAction::Continue
                        }
                        RecoveryAction::Backoff => {
                            consecutive_failures = consecutive_failures.saturating_add(1);
                            LoopAction::Sleep(backoff.calculate(consecutive_failures))
                        }
                        RecoveryAction::Exit => LoopAction::Exit,
                    }
                }
                crate::critical::TxOutcome::UnexpectedPanic => {
                    let recovery = handler.on_unexpected_panic();
                    match recovery {
                        RecoveryAction::Continue => {
                            consecutive_failures = 0;
                            LoopAction::Continue
                        }
                        RecoveryAction::Backoff => {
                            consecutive_failures = consecutive_failures.saturating_add(1);
                            LoopAction::Sleep(backoff.calculate(consecutive_failures))
                        }
                        RecoveryAction::Exit => LoopAction::Exit,
                    }
                }
            };

            match action {
                LoopAction::Continue => {}
                LoopAction::Sleep(d) => {
                    tokio::select! {
                        _ = tokio::time::sleep(d) => {}
                        _ = cancel.cancelled() => break,
                    }
                }
                LoopAction::Exit => break,
            }
        }
    })
}

// ---------------------------------------------------------------------------
// LoopRegistry — centralized loop lifecycle management
// ---------------------------------------------------------------------------

/// Centralized registry for managing loop handler lifecycles.
///
/// All loop tasks are started via `register` and tracked by name.
/// Supports dynamic registration (e.g. hot-scan new instances).
pub struct LoopRegistry {
    handles: HashMap<String, LoopHandle>,
}

impl LoopRegistry {
    /// Create an empty registry.
    pub fn new() -> Self {
        Self { handles: HashMap::new() }
    }

    /// Register and start a loop handler. If a handler with the same name
    /// already exists and is still running, it is cancelled first.
    pub fn register<H: LoopHandler>(&mut self, name: impl Into<String>, handler: H) {
        let name = name.into();
        if let Some(existing) = self.handles.get(&name) {
            if !existing.is_finished() {
                existing.cancel();
            }
        }
        let handle = run_handler_loop(name.clone(), handler, None);
        self.handles.insert(name, handle);
    }

    /// Remove entries for loops that have already exited.
    pub fn cleanup_finished(&mut self) {
        self.handles.retain(|_, h| !h.is_finished());
    }

    /// Check whether a named loop is registered and still running.
    pub fn is_running(&self, name: &str) -> bool {
        self.handles.get(name).map_or(false, |h| !h.is_finished())
    }

    /// Trigger cancellation for all registered loops.
    pub fn shutdown_all(&self) {
        for handle in self.handles.values() {
            handle.cancel();
        }
    }

    /// Wait for all registered loops to exit. Consumes the registry.
    pub async fn join_all(self) {
        for (_, handle) in self.handles {
            handle.join().await;
        }
    }

    /// Register and start a cron handler with periodic execution.
    pub fn register_cron<C: CronHandler>(&mut self, name: impl Into<String>, handler: C) {
        let name = name.into();
        if let Some(existing) = self.handles.get(&name) {
            if !existing.is_finished() {
                existing.cancel();
            }
        }
        let handle = run_cron_loop(name.clone(), handler, None);
        self.handles.insert(name, handle);
    }
}

// ---------------------------------------------------------------------------
// CronHandler — declarative cron/scheduled task trait
// ---------------------------------------------------------------------------

/// Declarative trait for periodic tasks. The framework manages:
/// - Timer with configurable interval
/// - `TxScope` creation per execution
/// - `catch_unwind` for panic recovery
/// - Cancel-aware sleep between executions
///
/// Simpler than `LoopHandler` — no skip/backoff logic needed since
/// the interval already provides natural spacing.
#[async_trait]
pub trait CronHandler: Send + 'static {
    /// The transaction type for this cron task's scope.
    type Tx: Transaction;

    /// Execute one cycle. Called with a fresh TxScope each time.
    async fn execute(&mut self, scope: &TxScope<Self::Tx>);

    /// Called when execute panics with a known fatal token.
    fn on_fatal(&mut self, token: <Self::Tx as Transaction>::Fatal, detail: Sealed<String>);

    /// Called when execute panics without a known fatal token. Default: log only.
    fn on_unexpected_panic(&mut self) {}

    /// Interval between executions.
    fn interval(&self) -> Duration;
}

/// Internal: run a CronHandler inside a spawn_loop task.
#[doc(hidden)]
pub fn run_cron_loop<C: CronHandler>(name: String, mut handler: C, db: Option<crate::db::Db>) -> LoopHandle {
    spawn_loop(name, |cancel| async move {
        loop {
            // Sleep first (cron pattern: wait then execute)
            let interval = handler.interval();
            tokio::select! {
                _ = tokio::time::sleep(interval) => {}
                _ = cancel.cancelled() => break,
            }
            if cancel.get() { break; }

            let scope = match &db {
                Some(d) => TxScope::<C::Tx>::new_with_db(d.clone()),
                None => TxScope::<C::Tx>::new(),
            };
            let outcome = crate::critical::run_tx_scope(&scope, async {
                handler.execute(&scope).await;
            }).await;

            match outcome {
                crate::critical::TxOutcome::Completed(()) => {}
                crate::critical::TxOutcome::Aborted(token, detail) => {
                    handler.on_fatal(token, detail);
                }
                crate::critical::TxOutcome::UnexpectedPanic => {
                    handler.on_unexpected_panic();
                }
            }
        }
    })
}

// ---------------------------------------------------------------------------
// ItemResult — batch item processing result
// ---------------------------------------------------------------------------

/// Result of processing a single item in a batch loop.
pub enum ItemResult {
    /// Item processed successfully.
    Ok,
    /// Item skipped (e.g., instance is idle).
    Skip,
}

// ---------------------------------------------------------------------------
// BatchHandler — declarative batch processing trait
// ---------------------------------------------------------------------------

/// Declarative trait for batch processing loops. The framework manages:
/// - Periodic items refresh with configurable interval
/// - External notification for immediate refresh (`notify_refresh`)
/// - Per-item `TxScope` creation and `catch_unwind`
/// - Adaptive polling (sleep when items empty, wake on notify)
///
/// Engine implements this trait; the framework runs the loop.
#[async_trait]
pub trait BatchHandler: Send + 'static {
    /// The item type to process.
    type Item: Send + Sync + Clone + Eq + std::hash::Hash + std::fmt::Display + 'static;

    /// The transaction type for per-item scopes.
    type Tx: Transaction;

    /// Refresh the items list within a transaction scope.
    ///
    /// Called periodically (on `refresh_interval`) and on `notify_refresh()`.
    /// The scope provides access to IO operations (e.g. `scope.fs()`).
    /// Panic is caught by the framework — old cached items are preserved.
    async fn items(&self, scope: &TxScope<Self::Tx>) -> Vec<Self::Item>;

    /// Process a single item with an independent `TxScope`.
    ///
    /// Each call gets a fresh scope + `catch_unwind`. Panic does not
    /// affect other items in the current round.
    async fn handle(&self, scope: &TxScope<Self::Tx>, item: Self::Item) -> ItemResult;

    /// Notification: an item's `handle()` panicked with a known fatal token.
    ///
    /// The framework continues to the next item regardless.
    /// Default: no-op.
    fn on_fatal(&self, _token: <Self::Tx as Transaction>::Fatal, _detail: Sealed<String>) {}

    /// Notification: an item's `handle()` panicked without a fatal token.
    ///
    /// Default: no-op.
    fn on_unexpected_panic(&self) {}
}

// ---------------------------------------------------------------------------
// BatchLoopConfig — batch loop configuration
// ---------------------------------------------------------------------------

/// Configuration for a batch processing loop.
pub struct BatchLoopConfig {
    /// How often to refresh the items list.
    pub refresh_interval: Duration,
    /// Optional interval between processing rounds (cron-like).
    /// `None` = continuous processing (next round starts after refresh wait).
    pub round_interval: Option<Duration>,
}

// ---------------------------------------------------------------------------
// BatchLoopHandle — batch loop notification handle
// ---------------------------------------------------------------------------

/// Handle for triggering batch loop notifications.
///
/// Each batch loop has a dedicated handle for external interaction.
/// Currently supports `notify_refresh()` to trigger immediate items refresh.
/// Future: `notify_wake(item)` for per-item idle wakeup.
#[derive(Clone)]
pub struct BatchLoopHandle {
    refresh_notify: Arc<tokio::sync::Notify>,
}

impl BatchLoopHandle {
    /// Trigger an immediate items refresh.
    ///
    /// Wakes the batch loop if it is sleeping, causing it to call
    /// `items()` and update the cached items list.
    pub fn notify_refresh(&self) {
        self.refresh_notify.notify_one();
    }
}

// ---------------------------------------------------------------------------
// run_batch_loop — internal batch loop runner
// ---------------------------------------------------------------------------

/// Start a batch handler loop.
///
/// Returns `(LoopHandle, BatchLoopHandle)` — the loop lifecycle handle
/// and the notification handle for external interaction.
#[doc(hidden)]
/// Helper: refresh items list with a temporary TxScope.
///
/// Returns `Some(items)` on success, `None` if the refresh panicked or aborted.
/// Callers keep the old cache on `None`.
async fn refresh_items<H: BatchHandler + Sync>(handler: &H, db: &Option<crate::db::Db>) -> Option<Vec<H::Item>> {
    let scope = match db {
        Some(d) => TxScope::<H::Tx>::new_with_db(d.clone()),
        None => TxScope::<H::Tx>::new(),
    };
    match crate::critical::run_tx_scope(&scope, async {
        handler.items(&scope).await
    }).await {
        crate::critical::TxOutcome::Completed(items) => Some(items),
        _ => None,
    }
}

pub fn run_batch_loop<H: BatchHandler + Sync>(
    name: String,
    handler: H,
    config: BatchLoopConfig,
    db: Option<crate::db::Db>,
) -> (LoopHandle, BatchLoopHandle) {
    let refresh_notify = Arc::new(tokio::sync::Notify::new());
    let notify_inner = refresh_notify.clone();

    let handler = Arc::new(handler);

    let loop_handle = spawn_loop(name, |cancel| async move {
        let mut active_tasks: HashMap<H::Item, (tokio::task::JoinHandle<()>, CancelToken)> = HashMap::new();

        // Initial refresh
        if let Some(items) = refresh_items(&*handler, &db).await {
            // Spawn initial item loops
            for item in &items {
                let item_cancel = CancelToken::new();
                let handle = spawn_item_loop::<H>(
                    handler.clone(), db.clone(), item.clone(),
                    item_cancel.clone(), config.round_interval,
                );
                active_tasks.insert(item.clone(), (handle, item_cancel));
            }
        }

        loop {
            if cancel.get() {
                break;
            }

            // Wait for next refresh cycle or external notification
            tokio::select! {
                _ = tokio::time::sleep(config.refresh_interval) => {}
                _ = notify_inner.notified() => {}
                _ = cancel.cancelled() => break,
            }

            if cancel.get() {
                break;
            }

            // Clean up finished tasks
            active_tasks.retain(|_, (handle, _)| !handle.is_finished());

            // Refresh items and reconcile active tasks
            if let Some(new_items) = refresh_items(&*handler, &db).await {
                let new_set: HashSet<H::Item> = new_items.iter().cloned().collect();
                let old_set: HashSet<H::Item> = active_tasks.keys().cloned().collect();

                // Cancel tasks for disappeared items
                for removed in old_set.difference(&new_set) {
                    if let Some((_, item_cancel)) = active_tasks.remove(removed) {
                        item_cancel.trigger();
                    }
                }

                // Spawn tasks for new items
                for added in new_set.difference(&old_set) {
                    let item_cancel = CancelToken::new();
                    let handle = spawn_item_loop::<H>(
                        handler.clone(), db.clone(), added.clone(),
                        item_cancel.clone(), config.round_interval,
                    );
                    active_tasks.insert(added.clone(), (handle, item_cancel));
                }
            }
        }

        // Shutdown: cancel all item loops and wait
        let mut handles = Vec::new();
        for (_, (handle, item_cancel)) in active_tasks.drain() {
            item_cancel.trigger();
            handles.push(handle);
        }
        for handle in handles {
            let _ = handle.await;
        }
    });

    let batch_handle = BatchLoopHandle { refresh_notify };
    (loop_handle, batch_handle)
}

fn spawn_item_loop<H: BatchHandler + Sync + 'static>(
    handler: Arc<H>,
    db: Option<crate::db::Db>,
    item: H::Item,
    cancel: CancelToken,
    round_interval: Option<Duration>,
) -> tokio::task::JoinHandle<()> {
    tokio::task::spawn(async move {
        loop {
            if cancel.get() {
                break;
            }

            let scope = match &db {
                Some(d) => TxScope::<H::Tx>::new_with_db(d.clone()),
                None => TxScope::<H::Tx>::new(),
            };
            let item_clone = item.clone();
            let outcome = crate::critical::run_tx_scope(&scope, async {
                handler.handle(&scope, item_clone).await
            }).await;
            match outcome {
                crate::critical::TxOutcome::Completed(ItemResult::Ok)
                | crate::critical::TxOutcome::Completed(ItemResult::Skip) => {}
                crate::critical::TxOutcome::Aborted(token, detail) => {
                    handler.on_fatal(token, detail);
                }
                crate::critical::TxOutcome::UnexpectedPanic => {
                    handler.on_unexpected_panic();
                }
            }

            // Wait between rounds, interruptible by cancel
            match round_interval {
                Some(interval) => {
                    tokio::select! {
                        _ = tokio::time::sleep(interval) => {}
                        _ = cancel.cancelled() => break,
                    }
                }
                None => {
                    // Yield to scheduler to prevent starving other tasks
                    tokio::task::yield_now().await;
                }
            }
        }
    })
}


// ---------------------------------------------------------------------------
// LivenessToken — opaque tunnel connection liveness token
// ---------------------------------------------------------------------------

/// Opaque liveness token for a tunnel connection.
///
/// Created by the tunnel framework (`accept_tunnel_connection`), passed to
/// `ConnectionRegistry::insert`. Engine uses this to check / force-disconnect.
#[derive(Clone)]
pub struct LivenessToken(pub(crate) Arc<AtomicBool>);

impl LivenessToken {
    pub(crate) fn new(raw: Arc<AtomicBool>) -> Self {
        Self(raw)
    }

    /// Returns `true` if the connection is still alive.
    pub fn is_alive(&self) -> bool {
        self.0.load(Ordering::Relaxed)
    }

    /// Force-mark as disconnected (for cleanup of stale connections).
    pub fn force_disconnect(&self) {
        self.0.store(false, Ordering::Relaxed)
    }
}

// ---------------------------------------------------------------------------
// ---------------------------------------------------------------------------
// SwapResult — result of PeerNetwork state transition
// ---------------------------------------------------------------------------

/// Result of a state transition on [`PeerNetwork`].
#[derive(Debug, PartialEq)]
pub enum SwapResult {
    /// The network was previously off.
    WasOff,
    /// The network was previously in server mode.
    WasServer,
    /// The network was previously in client mode.
    WasClient,
}

// ---------------------------------------------------------------------------
// PeerNetwork<C> — server/client network state manager
// ---------------------------------------------------------------------------

enum NetworkInner<C> {
    Off,
    Server(ServerData<C>),
    Client(ClientData),
}

struct ServerData<C> {
    entries: HashMap<String, (C, LivenessToken)>,
}

struct ClientData {
    liveness: LivenessToken,
}

/// Tri-state network manager: Off / Server / Client.
///
/// Replaces the combination of `ModeManager` + `ConnectionRegistry`.
/// - **Server mode**: manages multiple inbound client connections.
/// - **Client mode**: holds a single outbound connection to a server.
///
/// State transitions return a [`SwapResult`] describing the previous state.
pub struct PeerNetwork<C: Send + Sync + 'static> {
    inner: AsyncLock<NetworkInner<C>>,
}

impl<C: Send + Sync + 'static + Clone> PeerNetwork<C> {
    /// Create a new `PeerNetwork` in the `Off` state.
    pub fn new() -> Self {
        Self {
            inner: AsyncLock::new(NetworkInner::Off),
        }
    }

    // ── State transitions ──────────────────────────────────────────────────

    /// Switch to server mode. Returns the previous state.
    pub async fn become_server(&self) -> SwapResult {
        let mut guard = self.inner.lock().await;
        let old = std::mem::replace(&mut *guard, NetworkInner::Server(ServerData {
            entries: HashMap::new(),
        }));
        match old {
            NetworkInner::Off => SwapResult::WasOff,
            NetworkInner::Server(_) => SwapResult::WasServer,
            NetworkInner::Client(_) => SwapResult::WasClient,
        }
    }

    /// Switch to client mode with the given liveness token. Returns the previous state.
    pub async fn become_client(&self, liveness: LivenessToken) -> SwapResult {
        let mut guard = self.inner.lock().await;
        let old = std::mem::replace(&mut *guard, NetworkInner::Client(ClientData { liveness }));
        match old {
            NetworkInner::Off => SwapResult::WasOff,
            NetworkInner::Server(_) => SwapResult::WasServer,
            NetworkInner::Client(_) => SwapResult::WasClient,
        }
    }

    /// Shut down — transition to `Off`. Returns the previous state.
    pub async fn shutdown(&self) -> SwapResult {
        let mut guard = self.inner.lock().await;
        let old = std::mem::replace(&mut *guard, NetworkInner::Off);
        match old {
            NetworkInner::Off => SwapResult::WasOff,
            NetworkInner::Server(_) => SwapResult::WasServer,
            NetworkInner::Client(_) => SwapResult::WasClient,
        }
    }

    // ── State queries ──────────────────────────────────────────────────────

    /// Returns `true` if currently off.
    pub async fn is_off(&self) -> bool {
        matches!(&*self.inner.lock().await, NetworkInner::Off)
    }

    /// Returns `true` if currently in server mode.
    pub async fn is_server(&self) -> bool {
        matches!(&*self.inner.lock().await, NetworkInner::Server(_))
    }

    /// Returns `true` if currently in client mode.
    pub async fn is_client(&self) -> bool {
        matches!(&*self.inner.lock().await, NetworkInner::Client(_))
    }

    // ── Server-mode operations ─────────────────────────────────────────────

    /// Register a new client connection. No-op if not in server mode.
    pub async fn accept_client(&self, id: String, info: C, liveness: LivenessToken) {
        let mut guard = self.inner.lock().await;
        if let NetworkInner::Server(data) = &mut *guard {
            data.entries.insert(id, (info, liveness));
        }
    }

    /// Remove a client by id. No-op if not in server mode.
    pub async fn remove_client(&self, id: &str) {
        let mut guard = self.inner.lock().await;
        if let NetworkInner::Server(data) = &mut *guard {
            data.entries.remove(id);
        }
    }

    /// Returns the count of currently alive client connections.
    /// Returns 0 if not in server mode.
    pub async fn client_count(&self) -> usize {
        let guard = self.inner.lock().await;
        if let NetworkInner::Server(data) = &*guard {
            data.entries.values().filter(|(_, t)| t.is_alive()).count()
        } else {
            0
        }
    }

    /// Find a connected client by id. Returns `None` if not found, disconnected, or not in server mode.
    pub async fn find_client(&self, id: &str) -> Option<C> {
        let guard = self.inner.lock().await;
        if let NetworkInner::Server(data) = &*guard {
            data.entries.get(id)
                .filter(|(_, t)| t.is_alive())
                .map(|(v, _)| v.clone())
        } else {
            None
        }
    }

    /// Get a snapshot of all alive client connections.
    /// Returns empty vec if not in server mode.
    pub async fn get_connected_clients(&self) -> Vec<(String, C)> {
        let guard = self.inner.lock().await;
        if let NetworkInner::Server(data) = &*guard {
            data.entries.iter()
                .filter(|(_, (_, t))| t.is_alive())
                .map(|(k, (v, _))| (k.clone(), v.clone()))
                .collect()
        } else {
            vec![]
        }
    }

    /// Execute a closure with the list of alive clients. Returns `None` if not in server mode.
    pub async fn with_clients<F, R>(&self, f: F) -> Option<R>
    where
        F: FnOnce(Vec<(String, C)>) -> R,
    {
        let guard = self.inner.lock().await;
        if let NetworkInner::Server(data) = &*guard {
            let connected: Vec<(String, C)> = data.entries.iter()
                .filter(|(_, (_, t))| t.is_alive())
                .map(|(k, (v, _))| (k.clone(), v.clone()))
                .collect();
            Some(f(connected))
        } else {
            None
        }
    }

    /// Remove all disconnected clients and return the removed entries.
    /// Returns empty vec if not in server mode.
    pub async fn retain_clients(&self) -> Vec<(String, C, LivenessToken)> {
        let mut guard = self.inner.lock().await;
        if let NetworkInner::Server(data) = &mut *guard {
            let mut removed = Vec::new();
            data.entries.retain(|k, (v, t)| {
                if t.is_alive() {
                    true
                } else {
                    removed.push((k.clone(), v.clone(), t.clone()));
                    false
                }
            });
            removed
        } else {
            vec![]
        }
    }

    // ── Client-mode operations ─────────────────────────────────────────────

    /// Returns `true` if the server connection is alive.
    /// Returns `false` if not in client mode or if disconnected.
    pub async fn connection_alive(&self) -> bool {
        let guard = self.inner.lock().await;
        if let NetworkInner::Client(data) = &*guard {
            data.liveness.is_alive()
        } else {
            false
        }
    }
}

impl<C: Send + Sync + 'static + Clone> Default for PeerNetwork<C> {
    fn default() -> Self {
        Self::new()
    }
}

// ---------------------------------------------------------------------------
// ConnectionRegistry<V> — keyed connection pool with liveness filtering
// ---------------------------------------------------------------------------

struct ConnEntry<V> {
    value: V,
    liveness: LivenessToken,
}

/// A keyed registry of tunnel connections with automatic liveness filtering.
///
/// Query methods (`find_connected`, `get_all_connected`) automatically filter
/// out disconnected entries. Engine does not need to manually check liveness.
pub struct ConnectionRegistry<V: Clone + Send + Sync + 'static> {
    inner: AsyncCell<HashMap<String, ConnEntry<V>>>,
}

impl<V: Clone + Send + Sync + 'static> ConnectionRegistry<V> {
    /// Create an empty registry.
    pub fn new() -> Self {
        Self {
            inner: AsyncCell::new(HashMap::new()),
        }
    }

    /// Insert or replace a connection. Returns the old value if the key existed.
    pub async fn insert(&self, key: String, value: V, liveness: LivenessToken) -> Option<V> {
        let mut guard = self.inner.write().await;
        guard
            .insert(key, ConnEntry { value, liveness })
            .map(|e| e.value)
    }

    /// Remove a connection by key. Returns the value if it existed.
    pub async fn remove(&self, key: &str) -> Option<V> {
        let mut guard = self.inner.write().await;
        guard.remove(key).map(|e| e.value)
    }

    /// Update the value for an existing key via a closure. Returns `false` if not found.
    pub async fn update<F: FnOnce(&mut V)>(&self, key: &str, f: F) -> bool {
        let mut guard = self.inner.write().await;
        if let Some(entry) = guard.get_mut(key) {
            f(&mut entry.value);
            true
        } else {
            false
        }
    }

    /// Retain entries matching the predicate. Returns removed entries (with their tokens).
    pub async fn retain<F: FnMut(&str, &V) -> bool>(
        &self,
        mut f: F,
    ) -> Vec<(String, V, LivenessToken)> {
        let mut guard = self.inner.write().await;
        let mut removed = Vec::new();
        guard.retain(|k, e| {
            if f(k.as_str(), &e.value) {
                true
            } else {
                removed.push((k.clone(), e.value.clone(), e.liveness.clone()));
                false
            }
        });
        removed
    }

    /// Find the first connected entry matching the predicate.
    /// Automatically filters out disconnected entries.
    pub async fn find_connected<F: Fn(&str, &V) -> bool>(
        &self,
        predicate: F,
    ) -> Option<(String, V)> {
        let guard = self.inner.read().await;
        for (k, e) in guard.iter() {
            if e.liveness.is_alive() && predicate(k.as_str(), &e.value) {
                return Some((k.clone(), e.value.clone()));
            }
        }
        None
    }

    /// Get all connected entries.
    pub async fn get_all_connected(&self) -> Vec<(String, V)> {
        let guard = self.inner.read().await;
        guard
            .iter()
            .filter(|(_, e)| e.liveness.is_alive())
            .map(|(k, e)| (k.clone(), e.value.clone()))
            .collect()
    }

    /// Total number of entries (including disconnected).
    pub async fn len(&self) -> usize {
        self.inner.read().await.len()
    }

    /// Returns `true` if the registry is empty.
    pub async fn is_empty(&self) -> bool {
        self.inner.read().await.is_empty()
    }
}

impl<V: Clone + Send + Sync + 'static> Default for ConnectionRegistry<V> {
    fn default() -> Self {
        Self::new()
    }
}

// ---------------------------------------------------------------------------
// SharedList<T> — async shared list with clone-on-read
// ---------------------------------------------------------------------------

/// A shared, async-safe list. Clones on read; replaces on write.
///
/// Suitable for slave-side shared instance lists where readers outnumber writers.
pub struct SharedList<T: Clone + Send + Sync + 'static> {
    inner: AsyncLock<Vec<T>>,
}

impl<T: Clone + Send + Sync + 'static> SharedList<T> {
    /// Create a new list with the given initial items.
    pub fn new(items: Vec<T>) -> Self {
        Self {
            inner: AsyncLock::new(items),
        }
    }

    /// Get a snapshot of the current list (clones the entire vec).
    pub async fn get(&self) -> Vec<T> {
        self.inner.lock().await.clone()
    }

    /// Replace the list with new items.
    pub async fn set(&self, items: Vec<T>) {
        *self.inner.lock().await = items;
    }

    /// Append an item to the list.
    pub async fn push(&self, item: T) {
        self.inner.lock().await.push(item);
    }

    /// Returns the number of items.
    pub async fn len(&self) -> usize {
        self.inner.lock().await.len()
    }

    /// Returns `true` if empty.
    pub async fn is_empty(&self) -> bool {
        self.inner.lock().await.is_empty()
    }
}

impl<T: Clone + Send + Sync + 'static> Default for SharedList<T> {
    fn default() -> Self {
        Self::new(Vec::new())
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    // --- AtomicFlag tests ---

    #[test]
    fn atomic_flag_basic() {
        let flag = AtomicFlag::new();
        assert!(!flag.get());
        flag.set();
        assert!(flag.get());
    }

    #[test]
    fn atomic_flag_take() {
        let flag = AtomicFlag::new();
        flag.set();
        assert!(flag.take());
        assert!(!flag.get());
    }

    #[test]
    fn atomic_flag_clone() {
        let flag = AtomicFlag::new();
        let flag2 = flag.clone();
        flag.set();
        assert!(flag2.get());
    }

    #[test]
    fn atomic_flag_from_raw() {
        let raw = Arc::new(AtomicBool::new(false));
        let flag = AtomicFlag::from_raw(raw.clone());
        raw.store(true, Ordering::SeqCst);
        assert!(flag.get());
    }

    // --- ConnectionCache tests ---

    #[test]
    fn connection_cache_basic() {
        let cache = ConnectionCache::<String>::new();
        assert!(cache.get("a").is_none());
        cache.insert("a", "hello".to_owned());
        assert_eq!(cache.get("a"), Some("hello".to_owned()));
        cache.remove("a");
        assert!(cache.get("a").is_none());
    }

    #[test]
    fn connection_cache_get_or_insert() {
        let cache = ConnectionCache::<i32>::new();
        let v = cache.get_or_insert("k", 42);
        assert_eq!(v, 42);
        let v2 = cache.get_or_insert("k", 99);
        assert_eq!(v2, 42);
    }

    #[test]
    fn connection_cache_clone() {
        let cache = ConnectionCache::<String>::new();
        let cache2 = cache.clone();
        cache.insert("x", "shared".to_owned());
        assert_eq!(cache2.get("x"), Some("shared".to_owned()));
    }

    // --- AsyncCell tests ---

    #[tokio::test]
    async fn async_cell_basic() {
        let cell = AsyncCell::new(42);
        assert_eq!(*cell.read().await, 42);
        *cell.write().await = 100;
        assert_eq!(*cell.read().await, 100);
    }

    // --- AsyncLock tests ---

    #[tokio::test]
    async fn async_lock_basic() {
        let lock = AsyncLock::new(vec![1, 2, 3]);
        {
            let mut guard = lock.lock().await;
            guard.push(4);
        }
        let guard = lock.lock().await;
        assert_eq!(*guard, vec![1, 2, 3, 4]);
    }

    // --- channel tests ---

    #[test]
    fn channel_basic() {
        let (tx, mut rx) = channel::unbounded_channel::<i32>();
        tx.send(42).unwrap();
        assert_eq!(rx.try_recv().unwrap(), 42);
    }

    // --- CancelToken tests ---

    #[test]
    fn cancel_token_initially_false() {
        let ct = CancelToken::new();
        assert!(!ct.get());
    }

    #[test]
    fn cancel_token_trigger() {
        let ct = CancelToken::new();
        ct.trigger();
        assert!(ct.get());
    }

    #[test]
    fn cancel_token_clone_shares_state() {
        let ct = CancelToken::new();
        let ct2 = ct.clone();
        ct.trigger();
        assert!(ct2.get());
    }

    #[tokio::test]
    async fn cancel_token_cancelled_returns_immediately_if_triggered() {
        let ct = CancelToken::new();
        ct.trigger();
        // Should return immediately, not hang
        ct.cancelled().await;
        assert!(ct.get());
    }

    #[tokio::test]
    async fn cancel_token_cancelled_wakes_on_trigger() {
        let ct = CancelToken::new();
        let ct2 = ct.clone();
        let handle = tokio::spawn(async move {
            ct2.cancelled().await;
            true
        });
        // Give the spawned task time to start waiting
        tokio::time::sleep(Duration::from_millis(50)).await;
        ct.trigger();
        let result = handle.await.unwrap();
        assert!(result);
    }

    // --- spawn_loop tests ---

    #[tokio::test]
    async fn spawn_loop_basic() {
        let handle = spawn_loop("test-loop", |cancel| async move {
            loop {
                if cancel.get() {
                    break;
                }
                tokio::time::sleep(Duration::from_millis(10)).await;
            }
        });
        assert!(!handle.is_finished());
        handle.cancel();
        handle.join().await;
        // If we get here, the loop exited gracefully
    }

    #[tokio::test]
    async fn spawn_loop_cancel_with_select() {
        let handle = spawn_loop("test-select", |cancel| async move {
            tokio::select! {
                _ = cancel.cancelled() => {}
                _ = tokio::time::sleep(Duration::from_secs(60)) => {
                    panic!("should have been cancelled");
                }
            }
        });
        tokio::time::sleep(Duration::from_millis(50)).await;
        handle.cancel();
        handle.join().await;
    }

    #[tokio::test]
    async fn spawn_loop_is_finished() {
        let handle = spawn_loop("test-finish", |_cancel| async move {
            // Exit immediately
        });
        handle.join().await;
        // Task finished without cancel
    }

    #[tokio::test]
    async fn spawn_loop_abort() {
        let handle = spawn_loop("test-abort", |_cancel| async move {
            tokio::time::sleep(Duration::from_secs(60)).await;
        });
        handle.abort();
        // Give time for abort to take effect
        tokio::time::sleep(Duration::from_millis(50)).await;
        assert!(handle.is_finished());
    }
    // --- shutdown_all_loops tests ---

    #[tokio::test]
    async fn shutdown_all_loops_triggers_all() {
        // Note: this test shares global registry with other tests.
        // We spawn loops and verify they get cancelled.
        let h1 = spawn_loop("shutdown-test-1", |cancel| async move {
            cancel.cancelled().await;
        });
        let h2 = spawn_loop("shutdown-test-2", |cancel| async move {
            cancel.cancelled().await;
        });
        tokio::time::sleep(Duration::from_millis(50)).await;
        shutdown_all_loops();
        h1.join().await;
        h2.join().await;
        // Both exited — shutdown worked
    }

    // --- BackoffConfig tests ---

    #[test]
    fn backoff_config_default() {
        let cfg = BackoffConfig::default();
        assert_eq!(cfg.base_secs, 2);
        assert_eq!(cfg.max_secs, 120);
    }

    #[test]
    fn backoff_config_calculation() {
        let cfg = BackoffConfig::new(2, 120);
        assert_eq!(cfg.calculate(0).as_secs(), 2);   // 2 * 2^0 = 2
        assert_eq!(cfg.calculate(1).as_secs(), 4);   // 2 * 2^1 = 4
        assert_eq!(cfg.calculate(2).as_secs(), 8);   // 2 * 2^2 = 8
        assert_eq!(cfg.calculate(5).as_secs(), 64);  // 2 * 2^5 = 64
        assert_eq!(cfg.calculate(6).as_secs(), 120); // 2 * 2^6 = 128, capped at 120
        assert_eq!(cfg.calculate(10).as_secs(), 120); // capped
        assert_eq!(cfg.calculate(63).as_secs(), 120); // huge, capped
    }

    #[test]
    fn backoff_config_overflow() {
        let cfg = BackoffConfig::new(2, 120);
        // consecutive_failures > 63 should not panic
        assert_eq!(cfg.calculate(64).as_secs(), 120);
        assert_eq!(cfg.calculate(100).as_secs(), 120);
    }

    // --- LoopHandler / LoopRegistry tests ---

    use crate::critical::{ConceptMapToken, FatalToken, SuppressibleToken, Transaction};

    #[derive(Debug, Clone, Copy)]
    enum TestFatal { Boom }

    impl ConceptMapToken for TestFatal {
        fn operation_name(&self) -> &'static str { "test_loop" }
    }
    impl FatalToken for TestFatal {}

    #[derive(Debug, Clone, Copy)]
    enum TestSuppressible {}
    impl ConceptMapToken for TestSuppressible {
        fn operation_name(&self) -> &'static str { match *self {} }
    }
    impl SuppressibleToken for TestSuppressible {}

    struct TestTx;
    impl Transaction for TestTx {
        type Fatal = TestFatal;
        type Suppressible = TestSuppressible;
        fn name() -> &'static str { "test_loop" }
    }

    struct CountingHandler {
        ticks: std::sync::Arc<std::sync::atomic::AtomicU32>,
        max_ticks: u32,
    }

    #[async_trait]
    impl LoopHandler for CountingHandler {
        type Tx = TestTx;

        async fn tick(&mut self, _scope: &TxScope<TestTx>) -> LoopAction {
            let prev = self.ticks.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            if prev + 1 >= self.max_ticks {
                LoopAction::Exit
            } else {
                LoopAction::Continue
            }
        }

        fn on_fatal(&mut self, _token: TestFatal, _detail: Sealed<String>) -> RecoveryAction {
            RecoveryAction::Backoff
        }
    }

    #[tokio::test]
    async fn loop_handler_runs_to_exit() {
        let ticks = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
        let handler = CountingHandler { ticks: ticks.clone(), max_ticks: 3 };
        let handle = run_handler_loop("test-handler".to_string(), handler, None);
        handle.join().await;
        assert_eq!(ticks.load(std::sync::atomic::Ordering::SeqCst), 3);
    }

    #[tokio::test]
    async fn loop_handler_cancel() {
        struct SleepyHandler;
        #[async_trait]
        impl LoopHandler for SleepyHandler {
            type Tx = TestTx;
            async fn tick(&mut self, _scope: &TxScope<TestTx>) -> LoopAction {
                LoopAction::Sleep(Duration::from_secs(3600))
            }
            fn on_fatal(&mut self, _t: TestFatal, _d: Sealed<String>) -> RecoveryAction {
                RecoveryAction::Backoff
            }
        }
        let handle = run_handler_loop("test-sleepy".to_string(), SleepyHandler, None);
        tokio::time::sleep(Duration::from_millis(50)).await;
        handle.cancel();
        handle.join().await;
    }

    #[tokio::test]
    async fn loop_registry_register_and_finish() {
        let ticks = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
        let mut registry = LoopRegistry::new();
        registry.register("counter", CountingHandler { ticks: ticks.clone(), max_ticks: 2 });
        assert!(registry.is_running("counter"));
        assert!(!registry.is_running("nonexistent"));
        // Wait for handler to finish
        tokio::time::sleep(Duration::from_millis(100)).await;
        assert!(!registry.is_running("counter"));
        registry.cleanup_finished();
        assert_eq!(registry.handles.len(), 0);
    }

    #[tokio::test]
    async fn loop_registry_shutdown_and_join() {
        struct ForeverHandler;
        #[async_trait]
        impl LoopHandler for ForeverHandler {
            type Tx = TestTx;
            async fn tick(&mut self, _scope: &TxScope<TestTx>) -> LoopAction {
                LoopAction::Sleep(Duration::from_secs(3600))
            }
            fn on_fatal(&mut self, _t: TestFatal, _d: Sealed<String>) -> RecoveryAction {
                RecoveryAction::Backoff
            }
        }
        let mut registry = LoopRegistry::new();
        registry.register("forever-1", ForeverHandler);
        registry.register("forever-2", ForeverHandler);
        assert!(registry.is_running("forever-1"));
        assert!(registry.is_running("forever-2"));
        registry.shutdown_all();
        registry.join_all().await;
    }

    #[tokio::test]
    async fn loop_registry_re_register() {
        let ticks1 = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
        let ticks2 = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
        let mut registry = LoopRegistry::new();
        // Register a long-running handler
        struct WaitHandler { ticks: std::sync::Arc<std::sync::atomic::AtomicU32> }
        #[async_trait]
        impl LoopHandler for WaitHandler {
            type Tx = TestTx;
            async fn tick(&mut self, _scope: &TxScope<TestTx>) -> LoopAction {
                self.ticks.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
                LoopAction::Sleep(Duration::from_secs(3600))
            }
            fn on_fatal(&mut self, _t: TestFatal, _d: Sealed<String>) -> RecoveryAction {
                RecoveryAction::Backoff
            }
        }
        registry.register("worker", WaitHandler { ticks: ticks1.clone() });
        tokio::time::sleep(Duration::from_millis(50)).await;
        assert!(ticks1.load(std::sync::atomic::Ordering::SeqCst) >= 1);
        // Re-register with a different handler (cancels old one)
        registry.register("worker", WaitHandler { ticks: ticks2.clone() });
        tokio::time::sleep(Duration::from_millis(50)).await;
        assert!(ticks2.load(std::sync::atomic::Ordering::SeqCst) >= 1);
        registry.shutdown_all();
        registry.join_all().await;
    }
    // --- CronHandler tests ---

    struct CountingCron {
        count: std::sync::Arc<std::sync::atomic::AtomicU32>,
        interval_ms: u64,
    }

    #[async_trait]
    impl CronHandler for CountingCron {
        type Tx = TestTx;

        async fn execute(&mut self, _scope: &TxScope<TestTx>) {
            self.count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        }

        fn on_fatal(&mut self, _token: TestFatal, _detail: Sealed<String>) {}

        fn interval(&self) -> Duration {
            Duration::from_millis(self.interval_ms)
        }
    }

    #[tokio::test]
    async fn cron_handler_executes_periodically() {
        let count = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
        let handle = run_cron_loop("test-cron".to_string(), CountingCron {
            count: count.clone(),
            interval_ms: 30,
        }, None);
        // Wait enough time for a few executions
        tokio::time::sleep(Duration::from_millis(120)).await;
        handle.cancel();
        handle.join().await;
        let final_count = count.load(std::sync::atomic::Ordering::SeqCst);
        assert!(final_count >= 2, "expected at least 2 executions, got {}", final_count);
    }

    #[tokio::test]
    async fn cron_handler_cancel_during_sleep() {
        let count = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
        let handle = run_cron_loop("test-cron-cancel".to_string(), CountingCron {
            count: count.clone(),
            interval_ms: 5000, // long interval
        }, None);
        tokio::time::sleep(Duration::from_millis(50)).await;
        handle.cancel();
        handle.join().await;
        // Should have 0 executions (cancelled during first sleep)
        assert_eq!(count.load(std::sync::atomic::Ordering::SeqCst), 0);
    }

    #[tokio::test]
    async fn loop_registry_register_cron() {
        let count = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0));
        let mut registry = LoopRegistry::new();
        registry.register_cron("my-cron", CountingCron {
            count: count.clone(),
            interval_ms: 20,
        });
        assert!(registry.is_running("my-cron"));
        tokio::time::sleep(Duration::from_millis(200)).await;
        let final_count = count.load(std::sync::atomic::Ordering::SeqCst);
        assert!(final_count >= 2, "expected at least 2 cron executions, got {}", final_count);
        registry.shutdown_all();
        registry.join_all().await;
    }


    // -----------------------------------------------------------------------
    // BatchHandler tests
    // -----------------------------------------------------------------------

    use std::sync::atomic::AtomicU32;

    #[derive(Debug, Clone, Copy, PartialEq)]
    enum TestBatchFatal {}
    impl crate::critical::FatalToken for TestBatchFatal {}

    #[derive(Debug, Clone, Copy)]
    enum TestBatchSuppressible {}
    impl crate::critical::ConceptMapToken for TestBatchSuppressible {
        fn operation_name(&self) -> &'static str { match *self {} }
    }
    impl crate::critical::SuppressibleToken for TestBatchSuppressible {}

    struct TestBatchTx;
    impl crate::critical::Transaction for TestBatchTx {
        type Fatal = TestBatchFatal;
        type Suppressible = TestBatchSuppressible;
        fn name() -> &'static str { "test_batch" }
    }
    impl crate::critical::ConceptMapToken for TestBatchFatal {
        fn operation_name(&self) -> &'static str { "test_batch" }
    }

    struct TestBatchHandler {
        items_list: Arc<std::sync::Mutex<Vec<u32>>>,
        processed: Arc<AtomicU32>,
    }

    #[async_trait]
    impl BatchHandler for TestBatchHandler {
        type Item = u32;
        type Tx = TestBatchTx;

        async fn items(&self, _scope: &TxScope<Self::Tx>) -> Vec<u32> {
            self.items_list.lock().unwrap().clone()
        }

        async fn handle(&self, _scope: &TxScope<Self::Tx>, _item: u32) -> ItemResult {
            self.processed.fetch_add(1, Ordering::SeqCst);
            ItemResult::Ok
        }
    }

    #[tokio::test]
    async fn batch_handler_processes_items() {
        let items = Arc::new(std::sync::Mutex::new(vec![1u32, 2, 3]));
        let processed = Arc::new(AtomicU32::new(0));

        let handler = TestBatchHandler {
            items_list: items.clone(),
            processed: processed.clone(),
        };

        let config = BatchLoopConfig {
            refresh_interval: Duration::from_millis(50),
            round_interval: None,
        };

        let (loop_handle, _batch_handle) = run_batch_loop(
            "test-batch".to_string(),
            handler,
            config,
            None,
        );

        // Wait for at least one round
        tokio::time::sleep(Duration::from_millis(200)).await;

        // Should have processed items
        assert!(processed.load(Ordering::SeqCst) >= 3, "Expected >= 3 items processed");

        loop_handle.cancel();
        loop_handle.join().await;
    }

    #[tokio::test]
    async fn batch_handler_notify_refresh() {
        let items = Arc::new(std::sync::Mutex::new(Vec::<u32>::new()));
        let processed = Arc::new(AtomicU32::new(0));

        let handler = TestBatchHandler {
            items_list: items.clone(),
            processed: processed.clone(),
        };

        let config = BatchLoopConfig {
            refresh_interval: Duration::from_secs(60), // long interval
            round_interval: None,
        };

        let (loop_handle, batch_handle) = run_batch_loop(
            "test-batch-notify".to_string(),
            handler,
            config,
            None,
        );

        // Wait a bit, should process nothing (empty items)
        tokio::time::sleep(Duration::from_millis(100)).await;
        assert_eq!(processed.load(Ordering::SeqCst), 0);

        // Add items and notify
        *items.lock().unwrap() = vec![10, 20];
        batch_handle.notify_refresh();

        // Wait for refresh + processing
        tokio::time::sleep(Duration::from_millis(200)).await;
        assert!(processed.load(Ordering::SeqCst) >= 2, "Expected >= 2 after notify");

        loop_handle.cancel();
        loop_handle.join().await;
    }

    #[tokio::test]
    async fn batch_handler_cancel() {
        let items = Arc::new(std::sync::Mutex::new(vec![1u32]));
        let processed = Arc::new(AtomicU32::new(0));

        let handler = TestBatchHandler {
            items_list: items.clone(),
            processed: processed.clone(),
        };

        let config = BatchLoopConfig {
            refresh_interval: Duration::from_millis(50),
            round_interval: None,
        };

        let (loop_handle, _batch_handle) = run_batch_loop(
            "test-batch-cancel".to_string(),
            handler,
            config,
            None,
        );

        tokio::time::sleep(Duration::from_millis(100)).await;
        loop_handle.cancel();
        loop_handle.join().await;
    }

    #[tokio::test]
    async fn batch_handler_item_panic_continues() {
        #[derive(Debug, Clone, Copy, PartialEq)]
        enum PanicFatal {}
        impl crate::critical::FatalToken for PanicFatal {}

        #[derive(Debug, Clone, Copy)]
        enum PanicSuppressible {}
        impl crate::critical::ConceptMapToken for PanicSuppressible {
            fn operation_name(&self) -> &'static str { match *self {} }
        }
        impl crate::critical::SuppressibleToken for PanicSuppressible {}

        struct PanicTx;
        impl crate::critical::Transaction for PanicTx {
            type Fatal = PanicFatal;
            type Suppressible = PanicSuppressible;
            fn name() -> &'static str { "panic_batch" }
        }
        impl crate::critical::ConceptMapToken for PanicFatal {
            fn operation_name(&self) -> &'static str { "test_panic_batch" }
        }

        struct PanicBatchHandler {
            processed: Arc<AtomicU32>,
        }

        #[async_trait]
        impl BatchHandler for PanicBatchHandler {
            type Item = u32;
            type Tx = PanicTx;

            async fn items(&self, _scope: &TxScope<Self::Tx>) -> Vec<u32> {
                vec![1, 2, 3]
            }

            async fn handle(&self, _scope: &TxScope<Self::Tx>, item: u32) -> ItemResult {
                if item == 2 {
                    panic!("item 2 panics");
                }
                self.processed.fetch_add(1, Ordering::SeqCst);
                ItemResult::Ok
            }
        }

        let processed = Arc::new(AtomicU32::new(0));
        let handler = PanicBatchHandler {
            processed: processed.clone(),
        };

        let config = BatchLoopConfig {
            refresh_interval: Duration::from_millis(50),
            round_interval: Some(Duration::from_millis(50)),
        };

        let (loop_handle, _) = run_batch_loop(
            "test-batch-panic".to_string(),
            handler,
            config,
            None,
        );

        // Wait for at least one round
        tokio::time::sleep(Duration::from_millis(300)).await;

        // Items 1 and 3 should be processed (item 2 panics but doesn't stop others)
        assert!(processed.load(Ordering::SeqCst) >= 2, "Should process items 1 and 3 despite item 2 panic");

        loop_handle.cancel();
        loop_handle.join().await;
    }

    // ── LivenessToken ──────────────────────────────────────────────────────

    #[test]
    fn liveness_token_is_alive_and_force_disconnect() {
        let raw = Arc::new(AtomicBool::new(true));
        let token = LivenessToken::new(raw.clone());
        assert!(token.is_alive());
        token.force_disconnect();
        assert!(!token.is_alive());
        assert!(!raw.load(Ordering::SeqCst));
    }


    // ── PeerNetwork ────────────────────────────────────────────────────────

    #[tokio::test]
    async fn peer_network_transitions() {
        let net: PeerNetwork<u32> = PeerNetwork::new();
        assert!(net.is_off().await);
        assert!(!net.is_server().await);
        assert!(!net.is_client().await);

        // Off → Server
        let prev = net.become_server().await;
        assert_eq!(prev, SwapResult::WasOff);
        assert!(net.is_server().await);

        // Server → Client
        let raw = Arc::new(AtomicBool::new(true));
        let liveness = LivenessToken::new(raw.clone());
        let prev = net.become_client(liveness).await;
        assert_eq!(prev, SwapResult::WasServer);
        assert!(net.is_client().await);

        // Client → Off
        let prev = net.shutdown().await;
        assert_eq!(prev, SwapResult::WasClient);
        assert!(net.is_off().await);
    }

    #[tokio::test]
    async fn peer_network_server_operations() {
        let net: PeerNetwork<u32> = PeerNetwork::new();
        net.become_server().await;
        assert_eq!(net.client_count().await, 0);

        // accept two clients
        let raw_a = Arc::new(AtomicBool::new(true));
        let raw_b = Arc::new(AtomicBool::new(true));
        let tok_a = LivenessToken::new(raw_a.clone());
        let tok_b = LivenessToken::new(raw_b.clone());
        net.accept_client("alice".to_string(), 1u32, tok_a).await;
        net.accept_client("bob".to_string(), 2u32, tok_b).await;
        assert_eq!(net.client_count().await, 2);

        // find_client
        assert_eq!(net.find_client("alice").await, Some(1u32));
        assert_eq!(net.find_client("unknown").await, None);

        // disconnect bob → not visible
        raw_b.store(false, std::sync::atomic::Ordering::SeqCst);
        assert_eq!(net.client_count().await, 1);
        assert_eq!(net.find_client("bob").await, None);

        // get_connected_clients returns only alice
        let connected = net.get_connected_clients().await;
        assert_eq!(connected.len(), 1);
        assert_eq!(connected[0].0, "alice");

        // retain_clients removes bob
        let removed = net.retain_clients().await;
        assert_eq!(removed.len(), 1);
        assert_eq!(removed[0].0, "bob");
        assert_eq!(net.client_count().await, 1);

        // remove_client alice
        net.remove_client("alice").await;
        assert_eq!(net.client_count().await, 0);
    }

    #[tokio::test]
    async fn peer_network_client_operations() {
        let net: PeerNetwork<u32> = PeerNetwork::new();

        // Not in client mode → connection_alive returns false
        assert!(!net.connection_alive().await);

        // Switch to client mode
        let raw = Arc::new(AtomicBool::new(true));
        let liveness = LivenessToken::new(raw.clone());
        net.become_client(liveness).await;
        assert!(net.connection_alive().await);

        // Disconnect
        raw.store(false, std::sync::atomic::Ordering::SeqCst);
        assert!(!net.connection_alive().await);
    }

    #[tokio::test]
    async fn peer_network_with_clients() {
        let net: PeerNetwork<u32> = PeerNetwork::new();
        net.become_server().await;

        let raw = Arc::new(AtomicBool::new(true));
        let tok = LivenessToken::new(raw.clone());
        net.accept_client("x".to_string(), 42u32, tok).await;

        let result = net.with_clients(|clients| {
            clients.iter().map(|(_, v)| *v).sum::<u32>()
        }).await;
        assert_eq!(result, Some(42));

        // Returns None if not in server mode
        let net2: PeerNetwork<u32> = PeerNetwork::new();
        let result2 = net2.with_clients(|_| 0u32).await;
        assert!(result2.is_none());
    }

    // ── ConnectionRegistry ─────────────────────────────────────────────────

    #[tokio::test]
    async fn connection_registry_basic() {
        let reg: ConnectionRegistry<u32> = ConnectionRegistry::new();
        assert_eq!(reg.len().await, 0);
        assert!(reg.is_empty().await);

        let raw = Arc::new(AtomicBool::new(true));
        let token = LivenessToken::new(raw.clone());
        let old = reg.insert("a".to_string(), 1u32, token).await;
        assert!(old.is_none());
        assert_eq!(reg.len().await, 1);

        // find_connected returns connected entries
        let found = reg.find_connected(|k, _v| k == "a").await;
        assert!(found.is_some());
        let (key, val) = found.unwrap();
        assert_eq!(key, "a");
        assert_eq!(val, 1u32);

        // mark disconnected → not returned by find_connected
        raw.store(false, Ordering::SeqCst);
        let found = reg.find_connected(|k, _v| k == "a").await;
        assert!(found.is_none());

        // still in registry (len counts all)
        assert_eq!(reg.len().await, 1);

        // remove
        let removed = reg.remove("a").await;
        assert_eq!(removed, Some(1u32));
        assert_eq!(reg.len().await, 0);
    }

    #[tokio::test]
    async fn connection_registry_retain() {
        let reg: ConnectionRegistry<u32> = ConnectionRegistry::new();

        let t1 = LivenessToken::new(Arc::new(AtomicBool::new(true)));
        let t2 = LivenessToken::new(Arc::new(AtomicBool::new(true)));
        reg.insert("a".to_string(), 1u32, t1).await;
        reg.insert("b".to_string(), 2u32, t2).await;

        // retain only "a" → "b" is removed
        let removed = reg.retain(|k, _v| k == "a").await;
        assert_eq!(removed.len(), 1);
        assert_eq!(removed[0].0, "b");
        assert_eq!(reg.len().await, 1);
    }

    // ── SharedList ─────────────────────────────────────────────────────────

    #[tokio::test]
    async fn shared_list_basic() {
        let list: SharedList<u32> = SharedList::new(vec![1, 2, 3]);
        assert_eq!(list.len().await, 3);

        let items = list.get().await;
        assert_eq!(items, vec![1u32, 2, 3]);

        list.push(4u32).await;
        assert_eq!(list.len().await, 4);

        list.set(vec![10u32, 20]).await;
        assert_eq!(list.len().await, 2);
        assert_eq!(list.get().await, vec![10u32, 20]);
    }

}