//! Routing learning state — the persisted half of AgentNet-style
//! self-organization ([arXiv:2504.00587], design in
//! `docs/proposals/agentnet-self-organization.md`).
//!
//! Capability-similarity routing (`declagents.route`) starts cold: it ranks
//! agents purely by embedding similarity between a need and each agent's static
//! capability surface. This module is the memory that lets routing *learn*:
//!
//! - **Per-agent success stats** — which agents actually complete routed work,
//!   as an EMA so recent behavior dominates. Feeds a success *prior* blended
//!   into the ranking score.
//! - **Directed agent→agent edge weights** `w(i,j)` — reinforced when agent `i`
//!   forwards work to agent `j` and it succeeds; pruned below a threshold. This
//!   is AgentNet's learned topology: `w' = α·w + (1−α)·S`.
//!
//! File-backed at `~/.car/routing.json` (`CAR_ROUTING_PATH` overrides for
//! tests), atomic write-through, mirroring [`crate::declarative::DeclRegistry`]
//! hygiene. Last-writer-wins like the sibling registries — outcome recording is
//! best-effort and must never fail the routed call.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;

/// EMA smoothing factor for success rate and edge weights (AgentNet's `α`).
/// Higher → more weight on history, slower to react.
pub const EMA_ALPHA: f32 = 0.7;

/// Default edge-prune threshold (AgentNet's `θ_w`). Edges that decay below this
/// carry no useful routing signal and are dropped.
pub const PRUNE_THRESHOLD: f32 = 0.1;

/// The success prior returned for an agent with no recorded history — neutral,
/// so a brand-new agent neither out-ranks nor is buried by proven ones on the
/// prior term alone.
pub const NEUTRAL_PRIOR: f32 = 0.5;

/// EMA smoothing for the learned capability vector (AgentNet's `β`). Higher →
/// the centroid moves slowly, dominated by the agent's history of successes.
pub const CAPABILITY_EMA_BETA: f32 = 0.8;

/// Per-agent routing outcome stats.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
pub struct AgentStats {
    pub successes: u64,
    pub failures: u64,
    /// EMA of per-run success (1.0 = completed, 0.0 = failed). Seeded to the
    /// first observation, then blended by [`EMA_ALPHA`].
    pub ema_success_rate: f32,
    /// Learned capability centroid — the EMA of the (query-side) embeddings of
    /// needs this agent has *succeeded* at (AgentNet's `c_i`). Empty until the
    /// first success; folded by [`CAPABILITY_EMA_BETA`]. Skipped on the wire
    /// when empty to keep `routing.json` lean for agents with no history.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub learned_vector: Vec<f32>,
}

impl AgentStats {
    fn total(&self) -> u64 {
        self.successes + self.failures
    }
}

/// The on-disk routing state and the read-only snapshot handed to callers — one
/// shape, so a snapshot is just a clone.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
pub struct RoutingSnapshot {
    /// agent id → outcome stats.
    #[serde(default)]
    pub agents: HashMap<String, AgentStats>,
    /// from-agent → (to-agent → directed edge weight).
    #[serde(default)]
    pub edges: HashMap<String, HashMap<String, f32>>,
}

impl RoutingSnapshot {
    /// Success prior for ranking — [`NEUTRAL_PRIOR`] when the agent has no
    /// recorded history, else its EMA success rate.
    pub fn success_prior(&self, agent: &str) -> f32 {
        match self.agents.get(agent) {
            Some(s) if s.total() > 0 => s.ema_success_rate,
            _ => NEUTRAL_PRIOR,
        }
    }

    /// The agent's learned capability centroid, if it has succeeded at least
    /// once. None until then (ranking falls back to cold-start similarity).
    pub fn learned_capability(&self, agent: &str) -> Option<&[f32]> {
        self.agents
            .get(agent)
            .map(|s| s.learned_vector.as_slice())
            .filter(|v| !v.is_empty())
    }

    /// Directed edge weight `from → to` — how successfully `from` has forwarded
    /// work to `to`. 0.0 when no such edge has been recorded.
    pub fn edge_weight(&self, from: &str, to: &str) -> f32 {
        self.edges
            .get(from)
            .and_then(|m| m.get(to))
            .copied()
            .unwrap_or(0.0)
    }

    /// The strongest learned forward target from `agent`, if any — the peer
    /// `agent` has most successfully delegated to. None when no edges exist.
    pub fn best_forward(&self, agent: &str) -> Option<(String, f32)> {
        self.edges
            .get(agent)?
            .iter()
            .max_by(|a, b| a.1.total_cmp(b.1))
            .map(|(to, w)| (to.clone(), *w))
    }
}

/// EMA update. The first sample seeds the value directly (no blend against a
/// meaningless zero prior); subsequent samples blend by [`EMA_ALPHA`].
fn ema(prev: f32, sample: f32, seeded: bool) -> f32 {
    if seeded {
        EMA_ALPHA * prev + (1.0 - EMA_ALPHA) * sample
    } else {
        sample
    }
}

/// EMA-fold `sample` into `prev` in place. An empty `prev` (first success) or a
/// dimension change (model swap) reseeds directly rather than blending across
/// mismatched spaces.
fn ema_vector(prev: &mut Vec<f32>, sample: &[f32], beta: f32) {
    if prev.len() != sample.len() {
        *prev = sample.to_vec();
    } else {
        for (p, s) in prev.iter_mut().zip(sample) {
            *p = beta * *p + (1.0 - beta) * *s;
        }
    }
}

/// Prune directed edges below `theta` from an in-memory state. Returns the
/// count removed; empty per-source maps are dropped too. Shared by
/// [`RoutingStore::record_edge`] (inline decay) and
/// [`RoutingStore::prune_below`] (maintenance) so the two can't diverge.
fn prune_edges_in_place(state: &mut RoutingSnapshot, theta: f32) -> usize {
    let mut pruned = 0usize;
    for inner in state.edges.values_mut() {
        let before = inner.len();
        inner.retain(|_, w| *w >= theta);
        pruned += before - inner.len();
    }
    state.edges.retain(|_, inner| !inner.is_empty());
    pruned
}

/// File-backed routing learning store. See module docs.
pub struct RoutingStore {
    path: PathBuf,
    /// Serializes read-modify-write within this process so concurrent routed
    /// runs don't clobber each other's counter/EMA updates (the file's own
    /// last-writer-wins only protects cross-*process* integrity, not the lost
    /// in-process increments). One store instance per daemon (held in a
    /// `OnceLock`), so this is the authoritative serialization point.
    lock: std::sync::Mutex<()>,
}

impl RoutingStore {
    fn new(path: PathBuf) -> Self {
        Self {
            path,
            lock: std::sync::Mutex::new(()),
        }
    }

    /// `~/.car/routing.json` (or `CAR_ROUTING_PATH`).
    pub fn user_default() -> Result<Self, String> {
        if let Some(p) = std::env::var_os("CAR_ROUTING_PATH") {
            return Ok(Self::new(PathBuf::from(p)));
        }
        let home = std::env::var_os("HOME")
            .or_else(|| std::env::var_os("USERPROFILE"))
            .ok_or("cannot resolve home directory (HOME/USERPROFILE unset)")?;
        Ok(Self::new(PathBuf::from(home).join(".car").join("routing.json")))
    }

    pub fn at(path: impl Into<PathBuf>) -> Self {
        Self::new(path.into())
    }

    /// Read the state. A *missing* file reads as empty (cold start, expected). A
    /// *present but corrupt* file also reads as empty so a routed call never
    /// breaks — but that silently discards learned topology, so it's logged.
    fn read(&self) -> RoutingSnapshot {
        let contents = match std::fs::read_to_string(&self.path) {
            Ok(c) => c,
            Err(_) => return RoutingSnapshot::default(),
        };
        match serde_json::from_str(&contents) {
            Ok(state) => state,
            Err(e) => {
                tracing::warn!(
                    path = %self.path.display(),
                    error = %e,
                    "routing.json failed to parse — resetting learned routing state to empty"
                );
                RoutingSnapshot::default()
            }
        }
    }

    fn write(&self, state: &RoutingSnapshot) -> Result<(), String> {
        if let Some(parent) = self.path.parent() {
            std::fs::create_dir_all(parent)
                .map_err(|e| format!("create {}: {e}", parent.display()))?;
        }
        let json = serde_json::to_string_pretty(state).map_err(|e| e.to_string())?;
        // Atomic: write a sibling temp then rename over the target.
        let tmp = self.path.with_extension("json.tmp");
        std::fs::write(&tmp, json).map_err(|e| format!("write {}: {e}", tmp.display()))?;
        std::fs::rename(&tmp, &self.path)
            .map_err(|e| format!("rename into {}: {e}", self.path.display()))
    }

    /// Read-only snapshot for ranking and observability.
    pub fn snapshot(&self) -> RoutingSnapshot {
        self.read()
    }

    /// Record a routed run's outcome for the chosen agent.
    pub fn record_outcome(&self, agent: &str, ok: bool) -> Result<(), String> {
        let _guard = self.lock.lock().unwrap_or_else(|e| e.into_inner());
        let mut state = self.read();
        let stats = state.agents.entry(agent.to_string()).or_default();
        let seeded = stats.total() > 0;
        stats.ema_success_rate = ema(stats.ema_success_rate, if ok { 1.0 } else { 0.0 }, seeded);
        if ok {
            stats.successes += 1;
        } else {
            stats.failures += 1;
        }
        self.write(&state)
    }

    /// Fold a succeeded need's embedding into the agent's learned capability
    /// centroid (AgentNet's `c_i` reinforcement). No-op for an empty or
    /// non-finite embedding.
    ///
    /// The finiteness guard is load-bearing: serde serializes a NaN/Inf `f32`
    /// as JSON `null`, which then fails to deserialize as `f32` on the next
    /// read — and [`Self::read`] treats a parse failure as "reset to empty",
    /// so one poisoned component would silently and permanently wipe the whole
    /// store (every centroid, count, and edge). Refuse the sample instead. EMA
    /// over finite bounded inputs stays finite, so guarding the input suffices.
    pub fn record_capability(&self, agent: &str, task_emb: &[f32]) -> Result<(), String> {
        if task_emb.is_empty() || task_emb.iter().any(|x| !x.is_finite()) {
            return Ok(());
        }
        let _guard = self.lock.lock().unwrap_or_else(|e| e.into_inner());
        let mut state = self.read();
        let stats = state.agents.entry(agent.to_string()).or_default();
        ema_vector(&mut stats.learned_vector, task_emb, CAPABILITY_EMA_BETA);
        self.write(&state)
    }

    /// Reinforce (or weaken) the directed forward edge `from → to` by the
    /// outcome of that delegation. EMA from an implicit zero, so a single
    /// success lifts the weight partway and repeated successes climb toward 1.
    /// Edges that decay below [`PRUNE_THRESHOLD`] are dropped in the same write,
    /// so the edge set stays bounded without a separate maintenance pass.
    pub fn record_edge(&self, from: &str, to: &str, ok: bool) -> Result<(), String> {
        let _guard = self.lock.lock().unwrap_or_else(|e| e.into_inner());
        let mut state = self.read();
        let w = state
            .edges
            .entry(from.to_string())
            .or_default()
            .entry(to.to_string())
            .or_insert(0.0);
        *w = EMA_ALPHA * *w + (1.0 - EMA_ALPHA) * if ok { 1.0 } else { 0.0 };
        prune_edges_in_place(&mut state, PRUNE_THRESHOLD);
        self.write(&state)
    }

    /// Drop directed edges whose weight is below `theta`. Returns the number
    /// pruned. Available as an explicit maintenance op (e.g. to prune with a
    /// stricter threshold than the inline decay); `record_edge` already prunes
    /// at [`PRUNE_THRESHOLD`] on every write.
    pub fn prune_below(&self, theta: f32) -> Result<usize, String> {
        let _guard = self.lock.lock().unwrap_or_else(|e| e.into_inner());
        let mut state = self.read();
        let pruned = prune_edges_in_place(&mut state, theta);
        if pruned > 0 {
            self.write(&state)?;
        }
        Ok(pruned)
    }
}

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

    fn temp_store() -> (tempfile::TempDir, RoutingStore) {
        let dir = tempfile::tempdir().unwrap();
        let store = RoutingStore::at(dir.path().join("routing.json"));
        (dir, store)
    }

    #[test]
    fn no_history_yields_neutral_prior() {
        let (_d, store) = temp_store();
        assert_eq!(store.snapshot().success_prior("ghost"), NEUTRAL_PRIOR);
    }

    #[test]
    fn first_outcome_seeds_ema_then_blends() {
        let (_d, store) = temp_store();
        // First success seeds ema to 1.0 (no blend against zero).
        store.record_outcome("a", true).unwrap();
        let snap = store.snapshot();
        assert_eq!(snap.success_prior("a"), 1.0);
        assert_eq!(snap.agents["a"].successes, 1);

        // A subsequent failure blends: 0.7*1.0 + 0.3*0.0 = 0.7.
        store.record_outcome("a", false).unwrap();
        let snap = store.snapshot();
        assert!((snap.success_prior("a") - 0.7).abs() < 1e-6);
        assert_eq!(snap.agents["a"].failures, 1);
    }

    #[test]
    fn weak_edge_auto_pruned_on_decay() {
        let (_d, store) = temp_store();
        // One success: w = 0.7*0 + 0.3*1 = 0.3, above the 0.1 threshold.
        store.record_edge("a", "b", true).unwrap();
        assert!((store.snapshot().best_forward("a").unwrap().1 - 0.3).abs() < 1e-6);

        // A weak edge that decays below θ is dropped inline by record_edge — no
        // separate prune pass needed.
        store.record_edge("a", "c", true).unwrap(); // c at 0.3
        store.record_edge("a", "c", false).unwrap(); // c -> 0.21
        store.record_edge("a", "c", false).unwrap(); // c -> 0.147
        store.record_edge("a", "c", false).unwrap(); // c -> 0.1029
        store.record_edge("a", "c", false).unwrap(); // c -> 0.072 < 0.1 → pruned
        let snap = store.snapshot();
        assert_eq!(snap.edge_weight("a", "c"), 0.0); // gone
        assert_eq!(snap.best_forward("a").unwrap().0, "b"); // b survives
    }

    #[test]
    fn prune_below_runs_as_maintenance_with_strict_theta() {
        let (_d, store) = temp_store();
        store.record_edge("a", "b", true).unwrap(); // 0.3, survives inline prune
        // A stricter manual threshold removes it.
        assert_eq!(store.prune_below(0.5).unwrap(), 1);
        assert_eq!(store.snapshot().best_forward("a"), None);
    }

    #[test]
    fn edge_weight_reads_directed_pair() {
        let (_d, store) = temp_store();
        store.record_edge("a", "b", true).unwrap();
        let snap = store.snapshot();
        assert!((snap.edge_weight("a", "b") - 0.3).abs() < 1e-6);
        assert_eq!(snap.edge_weight("b", "a"), 0.0); // directed: reverse is absent
        assert_eq!(snap.edge_weight("a", "ghost"), 0.0);
    }

    #[test]
    fn best_forward_picks_max_weight() {
        let (_d, store) = temp_store();
        store.record_edge("a", "weak", true).unwrap(); // 0.3
        for _ in 0..3 {
            store.record_edge("a", "strong", true).unwrap(); // climbs >0.3
        }
        assert_eq!(store.snapshot().best_forward("a").unwrap().0, "strong");
    }

    #[test]
    fn capability_seeds_then_ema_folds() {
        let (_d, store) = temp_store();
        assert_eq!(store.snapshot().learned_capability("a"), None);

        // First success seeds the centroid directly.
        store.record_capability("a", &[1.0, 0.0, 0.0]).unwrap();
        assert_eq!(store.snapshot().learned_capability("a"), Some(&[1.0, 0.0, 0.0][..]));

        // Second folds by β=0.8: 0.8*1 + 0.2*0 = 0.8 on axis 0, 0.2 on axis 1.
        store.record_capability("a", &[0.0, 1.0, 0.0]).unwrap();
        let v = store.snapshot();
        let c = v.learned_capability("a").unwrap();
        assert!((c[0] - 0.8).abs() < 1e-6);
        assert!((c[1] - 0.2).abs() < 1e-6);
    }

    #[test]
    fn capability_reseeds_on_dimension_change() {
        let (_d, store) = temp_store();
        store.record_capability("a", &[1.0, 2.0]).unwrap();
        // A different dimensionality (model swap) replaces rather than blends.
        store.record_capability("a", &[9.0, 9.0, 9.0]).unwrap();
        assert_eq!(store.snapshot().learned_capability("a"), Some(&[9.0, 9.0, 9.0][..]));
    }

    #[test]
    fn empty_capability_embedding_is_noop() {
        let (_d, store) = temp_store();
        store.record_capability("a", &[]).unwrap();
        assert_eq!(store.snapshot().learned_capability("a"), None);
    }

    #[test]
    fn non_finite_capability_is_refused_and_store_survives() {
        let (_d, store) = temp_store();
        // Seed a valid centroid + some stats so we can prove they survive.
        store.record_capability("a", &[1.0, 2.0, 3.0]).unwrap();
        store.record_outcome("a", true).unwrap();
        // A NaN/Inf sample must be refused — persisting it would write JSON
        // `null` and wipe the whole store on the next read.
        store.record_capability("a", &[f32::NAN, 0.0, 0.0]).unwrap();
        store.record_capability("a", &[f32::INFINITY, 0.0, 0.0]).unwrap();
        let snap = store.snapshot();
        // Store is still readable and the prior centroid/stats are intact.
        assert_eq!(snap.learned_capability("a"), Some(&[1.0, 2.0, 3.0][..]));
        assert_eq!(snap.agents["a"].successes, 1);
    }

    #[test]
    fn corrupt_file_reads_as_empty() {
        let (dir, store) = temp_store();
        std::fs::write(dir.path().join("routing.json"), b"{ not json").unwrap();
        assert!(store.snapshot().agents.is_empty());
    }
}
