//! Aggregator registry — manages aggregate definitions and state.
//!
//! When an event is dispatched through the engine, matching aggregators
//! apply it to state managed by the registry's internal DashMap.

use std::any::{Any, TypeId};
use std::sync::Arc;

use anyhow::Result;
use dashmap::DashMap;
use uuid::Uuid;

use crate::reactor_v3::extract_prefix;
use crate::types::{LogCursor, StreamVersion};
use crate::upcaster::UpcasterRegistry;

// ── Aggregate state snapshots ────────────────────────────────────
//
// The `Aggregate` marker + `Apply<F: Fact>` extension traits live in
// `crate::aggregate_v3` (the v0.4 home). This module focuses on the
// registry-side machinery — type-erased dispatch, state storage,
// rollback, snapshots.

/// Per-event `(prev, next)` aggregate snapshots produced by a single fold.
///
/// Captured as stack-local pairs around `apply_event`, so transition guards
/// can read the actual prev/next for *this* event — not the racing `:prev`
/// DashMap slot, which is overwritten by every subsequent fold in the same
/// Phase 1 batch. See `tests/fan_in_races.rs::transition_*`.
#[derive(Default)]
pub struct TransitionSnapshots {
    inner: std::collections::HashMap<
        String,
        (Arc<dyn Any + Send + Sync>, Arc<dyn Any + Send + Sync>),
    >,
}

impl TransitionSnapshots {
    /// Empty snapshot map — used for ephemeral events that don't fold.
    pub fn empty() -> Self {
        Self::default()
    }

    /// Look up `(prev, next)` for an aggregate `A` at the given id.
    ///
    /// Returns `None` if this event did not affect `A` for that id, which
    /// transition guards should treat as "no transition — don't fire."
    pub fn get_pair<A: crate::aggregate_v3::Aggregate>(&self, id: Uuid) -> Option<(&A, &A)> {
        let key = format!("{}:{}", A::NAME, id);
        let (pre, post) = self.inner.get(&key)?;
        let pre_a = (**pre).downcast_ref::<A>()?;
        let post_a = (**post).downcast_ref::<A>()?;
        Some((pre_a, post_a))
    }

    pub(crate) fn insert(
        &mut self,
        key: String,
        prev: Arc<dyn Any + Send + Sync>,
        next: Arc<dyn Any + Send + Sync>,
    ) {
        self.inner.insert(key, (prev, next));
    }

    /// Iterate `(key, prev, next)` for every fold captured.
    /// Used by `AggregatorRegistry::notify_observer` to drive the
    /// inspector's `aggregate_folded` hook.
    pub fn iter(
        &self,
    ) -> impl Iterator<
        Item = (&str, &Arc<dyn Any + Send + Sync>, &Arc<dyn Any + Send + Sync>),
    > {
        self.inner
            .iter()
            .map(|(k, (prev, next))| (k.as_str(), prev, next))
    }
}

// ── Aggregator (type-erased event→aggregate applier) ────────────────

/// A type-erased aggregator that maps an event to an aggregate and applies it.
///
/// Clone is cheap — every non-trivial field is `Arc<dyn Fn>`. Used by
/// the `EngineBuilder` to build per-runner registry copies.
#[derive(Clone)]
pub struct Aggregator {
    /// The event prefix for matching (e.g. "scrape", "order_placed").
    pub event_prefix: String,
    /// TypeId of the event for fast matching.
    pub event_type_id: TypeId,
    /// The aggregate type string.
    pub aggregate_type: String,
    /// Extract the aggregate ID from JSON payload (deserializes internally).
    json_extract_id: Arc<dyn Fn(&serde_json::Value) -> Option<Uuid> + Send + Sync>,
    /// Deserialize JSON and apply to a type-erased aggregate (&mut dyn Any = &mut A).
    apply_to: Arc<dyn Fn(&mut dyn Any, serde_json::Value) -> Result<()> + Send + Sync>,
    /// Clone a type-erased aggregate state.
    clone_state: Arc<dyn Fn(&dyn Any) -> Box<dyn Any + Send + Sync> + Send + Sync>,
    /// Create a default aggregate state.
    default_state: Arc<dyn Fn() -> Box<dyn Any + Send + Sync> + Send + Sync>,
    /// Serialize aggregate state to JSON (for durable runtimes).
    serialize_state: Arc<dyn Fn(&dyn Any) -> Result<serde_json::Value> + Send + Sync>,
    /// Deserialize aggregate state from JSON (for durable runtimes).
    deserialize_state:
        Arc<dyn Fn(serde_json::Value) -> Result<Box<dyn Any + Send + Sync>> + Send + Sync>,
}

impl Aggregator {
    /// Construct an Aggregator that folds facts of type `F` into
    /// aggregate state `A`. Consumers read the folded state via
    /// `ctx.aggregate::<A>(stream_id)` inside their reactor / projector
    /// body.
    ///
    /// Stream id comes from `Fact::stream_id`; aggregate type string
    /// comes from `A::NAME` — explicit, stable across refactorings,
    /// portable to disk (backend `aggregate_type` columns).
    ///
    /// Registration via [`crate::EngineBuilder::with_aggregators`].
    pub fn for_type<A, F>() -> Self
    where
        A: crate::aggregate_v3::Aggregate
            + crate::aggregate_v3::Apply<F>
            + Clone
            + serde::Serialize
            + serde::de::DeserializeOwned,
        F: crate::fact::Fact,
    {
        let event_prefix = <F as crate::fact::Fact>::CATEGORY.to_string();
        let event_type_id = TypeId::of::<F>();
        let aggregate_type = <A as crate::aggregate_v3::Aggregate>::NAME.to_string();

        Self {
            event_prefix,
            event_type_id,
            aggregate_type,
            json_extract_id: Arc::new(|payload: &serde_json::Value| -> Option<Uuid> {
                let fact: F = serde_json::from_value(payload.clone()).ok()?;
                Some(<F as crate::fact::Fact>::stream_id(&fact))
            }),
            apply_to: Arc::new(|state: &mut dyn Any, data: serde_json::Value| -> Result<()> {
                let state = state
                    .downcast_mut::<A>()
                    .ok_or_else(|| anyhow::anyhow!("aggregate type mismatch in apply_to"))?;
                let fact: F = serde_json::from_value(data)?;
                <A as crate::aggregate_v3::Apply<F>>::apply(state, &fact);
                Ok(())
            }),
            clone_state: Arc::new(|state: &dyn Any| -> Box<dyn Any + Send + Sync> {
                let s = state.downcast_ref::<A>().unwrap();
                Box::new(s.clone())
            }),
            default_state: Arc::new(|| -> Box<dyn Any + Send + Sync> { Box::new(A::default()) }),
            serialize_state: Arc::new(|state: &dyn Any| -> Result<serde_json::Value> {
                let s = state
                    .downcast_ref::<A>()
                    .ok_or_else(|| anyhow::anyhow!("aggregate type mismatch in serialize_state"))?;
                Ok(serde_json::to_value(s)?)
            }),
            deserialize_state: Arc::new(
                |value: serde_json::Value| -> Result<Box<dyn Any + Send + Sync>> {
                    let s: A = serde_json::from_value(value)?;
                    Ok(Box::new(s))
                },
            ),
        }
    }

    /// Extract the aggregate ID from a JSON event payload.
    pub fn extract_id_from_json(&self, payload: &serde_json::Value) -> Option<Uuid> {
        (self.json_extract_id)(payload)
    }

    /// Apply this event's JSON data to a type-erased aggregate state.
    pub fn apply_to(&self, state: &mut dyn Any, data: serde_json::Value) -> Result<()> {
        (self.apply_to)(state, data)
    }

    /// Clone a type-erased aggregate state.
    pub fn clone_state(&self, state: &dyn Any) -> Box<dyn Any + Send + Sync> {
        (self.clone_state)(state)
    }

    /// Create a default aggregate state.
    pub fn default_state(&self) -> Box<dyn Any + Send + Sync> {
        (self.default_state)()
    }

    /// Serialize aggregate state to JSON (for durable runtimes).
    pub fn serialize_state(&self, state: &dyn Any) -> Result<serde_json::Value> {
        (self.serialize_state)(state)
    }

    /// Deserialize aggregate state from JSON (for durable runtimes).
    pub fn deserialize_state(
        &self,
        value: serde_json::Value,
    ) -> Result<Box<dyn Any + Send + Sync>> {
        (self.deserialize_state)(value)
    }
}

// ── AggregatorRegistry ──────────────────────────────────────────────

/// State entry that pairs aggregate state with its stream version.
///
/// Version travels with state in a single DashMap entry to avoid
/// split-brain between separate maps.
#[derive(Clone)]
struct StateEntry {
    state: Arc<dyn Any + Send + Sync>,
    /// Stream version from the Store (ZERO = never persisted / unknown).
    version: StreamVersion,
    /// Version at which last snapshot was taken (ZERO = never).
    snapshot_at_version: StreamVersion,
}

/// Captured aggregator state for rollback after a failed event-processing attempt.
///
/// Produced by [`AggregatorRegistry::capture_for_rollback`] before
/// `apply_event` mutates state, consumed by
/// [`AggregatorRegistry::restore_state`] when the engine needs to undo the
/// mutation (e.g. projection failure → retry).
pub(crate) struct AggregatorRollback {
    entries: Vec<RollbackEntry>,
}

struct RollbackEntry {
    key: String,
    prev_key: String,
    key_entry: Option<StateEntry>,
    prev_entry: Option<StateEntry>,
}

/// Registry of aggregators with owned in-memory state.
///
/// Holds `Aggregator` definitions plus the current folded state for
/// each `(aggregate_type, aggregate_id)` key. State lives in memory
/// for the registry's lifetime — there is no built-in persistence.
///
/// The v0.2.x engine includes snapshot integration on top of this
/// registry (`replay_events_onto`, `set_state`, `get_version`,
/// `update_snapshot_at_version`) for long-lived aggregates that span
/// engine instances. The v0.3 runners deliberately do NOT use that
/// integration — see `docs/aggregate-state-scope.md` for the design
/// rationale and the conditions under which it should be added back.
pub struct AggregatorRegistry {
    aggregators: Vec<Aggregator>,
    state: DashMap<String, StateEntry>,
}

impl AggregatorRegistry {
    pub fn new() -> Self {
        Self {
            aggregators: Vec::new(),
            state: DashMap::new(),
        }
    }

    pub fn register(&mut self, aggregator: Aggregator) {
        self.aggregators.push(aggregator);
    }

    /// Find all aggregators that handle the given durable name.
    ///
    /// Extracts the prefix from the durable name and matches against
    /// registered aggregator prefixes.
    pub fn find_by_durable_name(&self, durable_name: &str) -> Vec<&Aggregator> {
        let prefix = extract_prefix(durable_name);
        self.aggregators
            .iter()
            .filter(|a| a.event_prefix == prefix)
            .collect()
    }

    pub fn is_empty(&self) -> bool {
        self.aggregators.is_empty()
    }

    /// Apply an event to all matching aggregators, using internal state.
    ///
    /// For each matching aggregator:
    /// 1. Read current state (or create default)
    /// 2. Clone current state → prev snapshot
    /// 3. Apply event to cloned current state
    /// 4. Increment version atomically with state update
    ///
    /// State is stored as concrete types via `Arc<dyn Any>` — zero serialization overhead.
    pub fn apply_event(
        &self,
        event_type: &str,
        payload: &serde_json::Value,
    ) -> TransitionSnapshots {
        let mut snapshots = TransitionSnapshots::empty();
        let prefix = extract_prefix(event_type);
        let matching: Vec<&Aggregator> = self
            .aggregators
            .iter()
            .filter(|a| a.event_prefix == prefix)
            .collect();

        for agg in matching {
            let aggregate_id = match agg.extract_id_from_json(payload) {
                Some(id) => id,
                None => continue,
            };

            let key = format!("{}:{}", agg.aggregate_type, aggregate_id);
            let prev_key = format!("{}:prev", key);

            // Read pre-fold state. If no entry exists, materialise the default
            // *now* so we have an Arc to capture as `prev`.
            let (pre_state, current_version, snapshot_at) = match self.state.get(&key) {
                Some(entry) => (
                    entry.state.clone(),
                    entry.version,
                    entry.snapshot_at_version,
                ),
                None => {
                    let default: Arc<dyn Any + Send + Sync> = Arc::from(agg.default_state());
                    self.state.insert(
                        key.clone(),
                        StateEntry {
                            state: default.clone(),
                            version: StreamVersion::ZERO,
                            snapshot_at_version: StreamVersion::ZERO,
                        },
                    );
                    (default, StreamVersion::ZERO, StreamVersion::ZERO)
                }
            };

            // Clone pre-fold for mutation, apply, freeze as Arc.
            let mut next_state = agg.clone_state(pre_state.as_ref());
            if let Err(e) = agg.apply_to(next_state.as_mut(), payload.clone()) {
                tracing::error!("Failed to apply event to aggregate {}: {}", key, e);
            }
            let post_state: Arc<dyn Any + Send + Sync> = Arc::from(next_state);

            // Update the registry:
            //   :prev slot — kept for backward compat with `get_transition_*`
            //   readers, but it is racy under fan-in. New code reads from the
            //   `TransitionSnapshots` returned by this function instead.
            self.state.insert(
                prev_key,
                StateEntry {
                    state: pre_state.clone(),
                    version: StreamVersion::ZERO,
                    snapshot_at_version: StreamVersion::ZERO,
                },
            );
            self.state.insert(
                key.clone(),
                StateEntry {
                    state: post_state.clone(),
                    version: StreamVersion::from_raw(current_version.raw() + 1),
                    snapshot_at_version: snapshot_at,
                },
            );

            // Capture per-event (pre, post) for transition guards.
            snapshots.insert(key, pre_state, post_state);
        }

        snapshots
    }

    /// Replay a sequence of persisted events to reconstruct aggregate state.
    ///
    /// Takes `(event_type, payload)` pairs (decoupled from `PersistedEvent`).
    /// Uses short name matching so persisted events (e.g. `"OrderPlaced"`)
    /// match aggregators registered with full type paths.
    ///
    /// Upcasters are applied to each event payload before deserialization,
    /// transforming old schemas to the current version.
    ///
    /// Returns `None` if no aggregators are registered for this aggregate type.
    pub fn replay_events(
        &self,
        aggregate_type: &str,
        events: &[(&str, &serde_json::Value)],
        upcasters: &UpcasterRegistry,
    ) -> Result<Option<Box<dyn Any + Send + Sync>>> {
        // Find any aggregator for this aggregate_type to get default_state
        let first = self
            .aggregators
            .iter()
            .find(|a| a.aggregate_type == aggregate_type);

        let first = match first {
            Some(agg) => agg,
            None => return Ok(None),
        };

        let mut state = first.default_state();

        for (event_type, payload) in events {
            // Apply upcasters before deserialization (schema_version=0 as default)
            let upcasted_payload = upcasters.upcast(event_type, 0, (*payload).clone())?;

            // Find aggregators where the prefix matches
            let prefix = extract_prefix(event_type);
            let matching: Vec<&Aggregator> = self
                .aggregators
                .iter()
                .filter(|a| {
                    a.aggregate_type == aggregate_type
                        && a.event_prefix == prefix
                })
                .collect();

            for agg in matching {
                agg.apply_to(state.as_mut(), upcasted_payload.clone())?;
            }
        }

        Ok(Some(state))
    }

    /// Get the (prev, next) transition for an aggregate from internal state.
    /// Returns `(A::default(), A::default())` if no state exists.
    pub fn get_transition<A>(&self, id: Uuid) -> (A, A)
    where
        A: crate::aggregate_v3::Aggregate + Clone,
    {
        let key = format!("{}:{}", A::NAME, id);
        let prev_key = format!("{}:prev", key);

        let next = self
            .state
            .get(&key)
            .and_then(|entry| entry.state.downcast_ref::<A>().cloned())
            .unwrap_or_default();

        let prev = self
            .state
            .get(&prev_key)
            .and_then(|entry| entry.state.downcast_ref::<A>().cloned())
            .unwrap_or_default();

        (prev, next)
    }

    /// Get the (prev, next) transition as `Arc<A>` — zero-clone read access.
    pub fn get_transition_arc<A>(&self, id: Uuid) -> (Arc<A>, Arc<A>)
    where
        A: crate::aggregate_v3::Aggregate,
    {
        let key = format!("{}:{}", A::NAME, id);
        let prev_key = format!("{}:prev", key);

        let next = self
            .state
            .get(&key)
            .and_then(|entry| entry.state.clone().downcast::<A>().ok())
            .unwrap_or_else(|| Arc::new(A::default()));

        let prev = self
            .state
            .get(&prev_key)
            .and_then(|entry| entry.state.clone().downcast::<A>().ok())
            .unwrap_or_else(|| Arc::new(A::default()));

        (prev, next)
    }

    /// Get the (prev, next) transition for a singleton aggregate (uses `Uuid::nil()`).
    pub fn get_singleton<A>(&self) -> (A, A)
    where
        A: crate::aggregate_v3::Aggregate + Clone,
    {
        self.get_transition::<A>(Uuid::nil())
    }

    /// Get the (prev, next) transition for a singleton aggregate as `Arc<A>`.
    pub fn get_singleton_arc<A>(&self) -> (Arc<A>, Arc<A>)
    where
        A: crate::aggregate_v3::Aggregate,
    {
        self.get_transition_arc::<A>(Uuid::nil())
    }

    // ── Store integration helpers ──────────────────────────────

    /// Check if the DashMap has state for a given aggregate key.
    pub fn has_state(&self, key: &str) -> bool {
        self.state.contains_key(key)
    }

    /// Inject hydrated state + version into the DashMap.
    ///
    /// Used during cold-start hydration from the Store.
    pub fn set_state(&self, key: &str, state: Arc<dyn Any + Send + Sync>, version: StreamVersion, snapshot_at_version: StreamVersion) {
        self.state.insert(key.to_string(), StateEntry { state, version, snapshot_at_version });
    }

    /// Read the stream version from the DashMap entry.
    ///
    /// Returns 0 if no state exists (consistent with "version 0 = empty stream").
    pub fn get_version(&self, key: &str) -> StreamVersion {
        self.state
            .get(key)
            .map(|entry| entry.version)
            .unwrap_or(StreamVersion::ZERO)
    }

    /// Read the snapshot_at_version from the DashMap entry.
    ///
    /// Returns 0 if no state exists (consistent with "never snapshotted").
    pub fn get_snapshot_at_version(&self, key: &str) -> StreamVersion {
        self.state
            .get(key)
            .map(|entry| entry.snapshot_at_version)
            .unwrap_or(StreamVersion::ZERO)
    }

    /// Update snapshot_at_version after saving a snapshot.
    pub fn update_snapshot_at_version(&self, key: &str, version: StreamVersion) {
        if let Some(mut entry) = self.state.get_mut(key) {
            entry.snapshot_at_version = version;
        }
    }

    /// Get a clone of the state for a given aggregate key.
    ///
    /// Returns `None` if no state exists. Used by `save_snapshot`.
    pub fn get_state(&self, key: &str) -> Option<Arc<dyn Any + Send + Sync>> {
        self.state.get(key).map(|entry| entry.state.clone())
    }

    /// Remove cached state for an aggregate, forcing re-hydration on next access.
    ///
    /// Used for multi-node sync: after ingesting foreign events, invalidate
    /// the aggregate so the next settle loop hydrates fresh from the Store.
    pub fn remove_state(&self, key: &str) {
        self.state.remove(key);
        self.state.remove(&format!("{}:prev", key));
    }

    /// Return all unique aggregate type names registered.
    pub fn unique_aggregate_types(&self) -> Vec<&str> {
        let mut seen = std::collections::HashSet::new();
        self.aggregators
            .iter()
            .map(|a| a.aggregate_type.as_str())
            .filter(|t| seen.insert(*t))
            .collect()
    }

    /// Find the first aggregator registered for a given aggregate type string.
    ///
    /// Used to access `deserialize_state` / `default_state` during hydration.
    pub fn find_first_by_aggregate_type(&self, aggregate_type: &str) -> Option<&Aggregator> {
        self.aggregators
            .iter()
            .find(|a| a.aggregate_type == aggregate_type)
    }

    /// Push each `(key, next)` from a fold's snapshots into the
    /// observer's `aggregate_folded` hook. Used by runners that
    /// folded an event and want to surface state-after-fold to
    /// inspector / telemetry.
    ///
    /// Resolves the matching `Aggregator` by `aggregate_type` (the
    /// key's prefix before `:`) to obtain a typed serializer; the
    /// observer receives JSON, not type-erased `Any`.
    pub fn notify_observer(
        &self,
        snapshots: &TransitionSnapshots,
        observer: &dyn crate::reactor_observer::ReactorObserver,
        correlation_id: Uuid,
        position: LogCursor,
        event_id: Uuid,
    ) {
        for (key, _prev, next) in snapshots.iter() {
            // Key format: "{aggregate_type}:{aggregate_id}". Split
            // once at the first ':'.
            let aggregate_type = key.split(':').next().unwrap_or("");
            let Some(agg) = self.find_first_by_aggregate_type(aggregate_type) else {
                continue;
            };
            match agg.serialize_state(next.as_ref()) {
                Ok(state_json) => observer.aggregate_folded(
                    correlation_id,
                    position,
                    event_id,
                    key,
                    state_json,
                ),
                Err(e) => tracing::warn!(
                    aggregate_key = %key,
                    error = %e,
                    "notify_observer: serialize_state failed; skipping snapshot"
                ),
            }
        }
    }

    /// Capture pre-mutation state for the aggregates that would be affected
    /// by `apply_event(event_type, payload)`. The returned handle can be
    /// passed to [`restore_state`](Self::restore_state) to undo the mutation
    /// — used by the engine to roll back aggregator state when
    /// `process_event_inner` fails after `apply_event` already ran.
    ///
    /// Captures the existing entry for both `key` and `key:prev` (or absence).
    pub(crate) fn capture_for_rollback(
        &self,
        event_type: &str,
        payload: &serde_json::Value,
    ) -> AggregatorRollback {
        let prefix = extract_prefix(event_type);
        let mut entries = Vec::new();
        for agg in self.aggregators.iter().filter(|a| a.event_prefix == prefix) {
            let aggregate_id = match agg.extract_id_from_json(payload) {
                Some(id) => id,
                None => continue,
            };
            let key = format!("{}:{}", agg.aggregate_type, aggregate_id);
            let prev_key = format!("{}:prev", key);
            let key_entry = self.state.get(&key).map(|e| e.clone());
            let prev_entry = self.state.get(&prev_key).map(|e| e.clone());
            entries.push(RollbackEntry {
                key,
                prev_key,
                key_entry,
                prev_entry,
            });
        }
        AggregatorRollback { entries }
    }

    /// Restore aggregator state captured by [`capture_for_rollback`](Self::capture_for_rollback).
    ///
    /// Each captured entry is restored to its prior value, or removed if it
    /// did not exist before. Idempotent if called twice with the same handle.
    pub(crate) fn restore_state(&self, rollback: AggregatorRollback) {
        for entry in rollback.entries {
            match entry.key_entry {
                Some(state_entry) => {
                    self.state.insert(entry.key, state_entry);
                }
                None => {
                    self.state.remove(&entry.key);
                }
            }
            match entry.prev_entry {
                Some(state_entry) => {
                    self.state.insert(entry.prev_key, state_entry);
                }
                None => {
                    self.state.remove(&entry.prev_key);
                }
            }
        }
    }

    /// Replay events onto an existing state (for snapshot + partial replay).
    ///
    /// Matches events by short type name (same as `replay_events`).
    /// Upcasters are applied before deserialization.
    pub fn replay_events_onto(
        &self,
        aggregate_type: &str,
        state: &mut dyn Any,
        events: &[(&str, &serde_json::Value)],
        upcasters: &UpcasterRegistry,
    ) -> Result<()> {
        for (event_type, payload) in events {
            let upcasted = upcasters.upcast(event_type, 0, (*payload).clone())?;

            let prefix = extract_prefix(event_type);
            let matching: Vec<&Aggregator> = self
                .aggregators
                .iter()
                .filter(|a| {
                    a.aggregate_type == aggregate_type
                        && a.event_prefix == prefix
                })
                .collect();

            for agg in matching {
                agg.apply_to(state, upcasted.clone())?;
            }
        }

        Ok(())
    }
}

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