armdb 0.4.1

use std::any::Any;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use armour_core::{
    CollectionKind, CollectionNode, FieldRef, GetRefs, GetType, OnMissing, RelationEdge,
    RelationKind, SchemaGraph, StorageClass, Typ,
};

use super::access::{BrandedKey, IndexTarget, SchemaCollection};
use super::hook_factory::IndexHook;
use crate::CollectionMeta;
use crate::key::Key;

type VisitIdsFn = Box<dyn Fn(&mut dyn FnMut(&[u8], RefIn, &'static str, u64)) + Send + Sync>;
type ScanAnyFn = Box<dyn Fn(&mut dyn FnMut(&[u8], &dyn Any, &dyn Any)) + Send + Sync>;
type ContainsAnyFn = Box<dyn Fn(&dyn Any) -> Option<bool> + Send + Sync>;
type WithValueAnyFn = Box<dyn Fn(&dyn Any, &mut dyn FnMut(&dyn Any)) -> Option<bool> + Send + Sync>;
type BrandedContainsFn = Box<dyn Fn(u64) -> bool + Send + Sync>;
type BrandedOption = Option<(&'static str, BrandedContainsFn)>;
type IndexProjectFn =
    Box<dyn Fn(&dyn Any, &dyn Any) -> Option<(Vec<u8>, Box<dyn Any>, Box<dyn Any>)> + Send + Sync>;
type ValueEqFn = Box<dyn Fn(&dyn Any, &dyn Any) -> Option<bool> + Send + Sync>;
type DenormProjectFn =
    Box<dyn Fn(&dyn Any, &dyn Any) -> Option<(Vec<u8>, Box<dyn Any>)> + Send + Sync>;
type DenormConsistentFn = Box<dyn Fn(&dyn Any, &dyn Any) -> Option<bool> + Send + Sync>;
type CounterReadFn = Box<dyn Fn(&dyn Any) -> Option<u64> + Send + Sync>;
type CounterPrefixFn = Box<dyn Fn(&dyn Any) -> Option<Vec<u8>> + Send + Sync>;

/// Registry consistency error reported by [`SchemaRegistry::finish`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum RegistryBuildError {
    #[error("collection `{0}` registered twice")]
    DuplicateCollection(&'static str),
    #[error("brand `{brand}` owned by two entities: `{first}` and `{second}`")]
    DuplicateEntityBrand {
        brand: &'static str,
        first: &'static str,
        second: &'static str,
    },
}

/// Display label for an FK edge in the graph/report (not extraction — closures do that).
#[derive(Debug, Clone)]
pub struct FkLabel(pub String);

impl FkLabel {
    pub fn field(name: &str) -> Self {
        Self(name.to_string())
    }
}

/// Where an auto-`FieldRef` was found.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(in crate::schema) enum RefIn {
    Key,
    Value,
}

pub(super) struct CollectionEntry {
    pub name: &'static str,
    pub kind: CollectionKind,
    pub typ: Typ,
    pub self_brand: Option<&'static str>,
    pub storage: StorageClass,
    pub auto_refs: Vec<(RefIn, FieldRef)>,
    /// Scan with visit_ids on key and value; callback: (key_bytes, key|value, brand, id).
    pub visit_ids: VisitIdsFn,
    /// Scan with erased types; callback: (key_bytes, &K as Any, &V as Any).
    pub scan_any: ScanAnyFn,
    /// None — wrong key type (declaration error), Some(b) — presence.
    pub contains_any: ContainsAnyFn,
    /// Some(false) — no record, Some(true) — record existed and `f` was called.
    pub with_value_any: WithValueAnyFn,
    /// Only for `register_entity` with `BrandedKey`.
    pub contains_branded: Option<BrandedContainsFn>,
}

/// Result of an explicit FK key-builder.
pub(super) enum FkBuild {
    /// No reference (e.g. `Option::None`) — valid.
    NoRef,
    Key(Box<dyn Any>),
    TypeMismatch,
}

type FkBuildFn = Box<dyn Fn(&dyn Any, &dyn Any) -> FkBuild + Send + Sync>;

pub(super) enum Decl {
    Index {
        source: &'static str,
        index: &'static str,
        /// (src key bytes, src K any, src V any) -> (idx key bytes, idx K any, idx V any)
        project: IndexProjectFn,
        /// compare expected vs actual index value
        value_eq: ValueEqFn,
    },
    Fk {
        from: &'static str,
        to: &'static str,
        label: String,
        on_missing: OnMissing,
        /// (src K any, src V any) -> target key; inner `None` = no reference (valid).
        build: FkBuildFn,
    },
    Denorm {
        source: &'static str,
        denorm: &'static str,
        project: DenormProjectFn,
        consistent: DenormConsistentFn,
    },
    Counter {
        collection: &'static str,
        field: &'static str,
        source: &'static str,
        read: CounterReadFn,
        prefix: CounterPrefixFn,
    },
}

#[derive(Default)]
pub struct SchemaRegistry {
    pub(super) entries: Vec<CollectionEntry>,
    pub(super) decls: Vec<Decl>,
    pub(super) external: Vec<(&'static str, CollectionKind)>,
}

/// Registry after a successful consistency check.
pub struct BuiltRegistry {
    pub(super) reg: SchemaRegistry,
    /// brand -> entity collection name.
    pub(super) brand_targets: HashMap<&'static str, &'static str>,
}

impl SchemaRegistry {
    pub fn new() -> Self {
        Self::default()
    }

    /// Register a collection with an open handle.
    pub fn register<S: SchemaCollection>(&mut self, handle: &Arc<S>, kind: CollectionKind) {
        self.register_inner(handle, kind, None);
    }

    /// Entity with branded `SelfId` — becomes the auto-FK target for its brand.
    pub fn register_entity<S: SchemaCollection>(&mut self, handle: &Arc<S>)
    where
        S::K: BrandedKey,
    {
        let h = handle.clone();
        let contains =
            Box::new(move |id: u64| h.contains_key(&S::K::from_id(id))) as BrandedContainsFn;
        self.register_inner(
            handle,
            CollectionKind::Entity,
            Some((S::K::BRAND, contains)),
        );
    }

    /// Graph node without a handle (no validation; clears UnclassifiedCollection).
    pub fn register_external(&mut self, name: &'static str, kind: CollectionKind) {
        self.external.push((name, kind));
    }

    fn register_inner<S: SchemaCollection>(
        &mut self,
        handle: &Arc<S>,
        kind: CollectionKind,
        branded: BrandedOption,
    ) {
        let mut auto_refs: Vec<(RefIn, FieldRef)> =
            S::K::refs().into_iter().map(|r| (RefIn::Key, r)).collect();
        auto_refs.extend(S::V::refs().into_iter().map(|r| (RefIn::Value, r)));

        let h = handle.clone();
        #[allow(clippy::type_complexity)]
        let visit_ids = Box::new(move |f: &mut dyn FnMut(&[u8], RefIn, &'static str, u64)| {
            h.scan(&mut |k, v| {
                let kb = k.as_bytes();
                k.visit_ids(&mut |brand, id| f(kb, RefIn::Key, brand, id));
                v.visit_ids(&mut |brand, id| f(kb, RefIn::Value, brand, id));
            });
        }) as VisitIdsFn;

        let h = handle.clone();
        let scan_any: ScanAnyFn = Box::new(move |f| {
            h.scan(&mut |k, v| f(k.as_bytes(), k as &dyn Any, v as &dyn Any));
        });

        let h = handle.clone();
        let contains_any =
            Box::new(move |key: &dyn Any| key.downcast_ref::<S::K>().map(|k| h.contains_key(k)))
                as ContainsAnyFn;

        let h = handle.clone();
        let with_value_any = Box::new(move |key: &dyn Any, f: &mut dyn FnMut(&dyn Any)| {
            key.downcast_ref::<S::K>()
                .map(|k| h.with_value(k, &mut |v| f(v as &dyn Any)))
        }) as WithValueAnyFn;

        let (self_brand, contains_branded) = match branded {
            // An unbranded ID (`Fuid<()>` / `Id64<()>`, brand "") is never an entity
            // brand target — consistent with the empty-brand-as-no-reference contract.
            // Drop the brand so it can neither collide as `DuplicateEntityBrand("")`
            // nor export an empty `self_brand` into the graph.
            Some(("", _)) => (None, None),
            Some((b, c)) => (Some(b), Some(c)),
            None => (None, None),
        };

        self.entries.push(CollectionEntry {
            name: S::V::NAME,
            kind,
            typ: <S::V as GetType>::TYPE,
            self_brand,
            storage: S::STORAGE,
            auto_refs,
            visit_ids,
            scan_any,
            contains_any,
            with_value_any,
            contains_branded,
        });
    }

    /// Consistency check: unique entity brands, unique names.
    /// Declarations are validated as they are added in later tasks.
    pub fn finish(self) -> Result<BuiltRegistry, RegistryBuildError> {
        let mut brand_targets: HashMap<&'static str, &'static str> = HashMap::new();
        let mut names: HashSet<&'static str> = HashSet::new();
        for e in &self.entries {
            if !names.insert(e.name) {
                return Err(RegistryBuildError::DuplicateCollection(e.name));
            }
            if e.kind == CollectionKind::Entity
                && let Some(brand) = e.self_brand
                && let Some(prev) = brand_targets.insert(brand, e.name)
            {
                return Err(RegistryBuildError::DuplicateEntityBrand {
                    brand,
                    first: prev,
                    second: e.name,
                });
            }
        }
        Ok(BuiltRegistry {
            reg: self,
            brand_targets,
        })
    }

    /// Declare a secondary index: returns a hook for opening the source collection.
    /// The index must be opened before the source (the hook captures its `Arc`).
    ///
    /// Several sources may feed one index collection (one declaration per
    /// source): the validator's orphan check runs against the union of
    /// expected keys from all declared sources of that index.
    pub fn secondary_index<SV, I, P>(
        &mut self,
        index: Arc<I>,
        project: P,
    ) -> IndexHook<SV::SelfId, SV, I, P>
    where
        SV: CollectionMeta + Send + Sync + 'static,
        SV::SelfId: Key + 'static,
        I: IndexTarget,
        I::V: PartialEq,
        P: Fn(&SV::SelfId, &SV) -> (I::K, I::V) + Clone + Send + Sync + 'static,
    {
        let p = project.clone();
        self.decls.push(Decl::Index {
            source: SV::NAME,
            index: <I::V as CollectionMeta>::NAME,
            project: Box::new(move |k, v| {
                let k = k.downcast_ref::<SV::SelfId>()?;
                let v = v.downcast_ref::<SV>()?;
                let (ik, iv) = p(k, v);
                Some((
                    ik.as_bytes().to_vec(),
                    Box::new(ik) as Box<dyn Any>,
                    Box::new(iv) as Box<dyn Any>,
                ))
            }),
            value_eq: Box::new(|expected, actual| {
                let e = expected.downcast_ref::<I::V>()?;
                let a = actual.downcast_ref::<I::V>()?;
                Some(e == a)
            }),
        });
        IndexHook {
            index,
            project,
            _pd: std::marker::PhantomData,
        }
    }

    /// Explicit FK: typed key-builder (covers unbranded fields and composite targets).
    pub fn fk<SV, TV, F>(&mut self, on_missing: OnMissing, label: FkLabel, build: F)
    where
        SV: CollectionMeta + 'static,
        SV::SelfId: Key + 'static,
        TV: CollectionMeta + 'static,
        TV::SelfId: Key + 'static,
        F: Fn(&SV::SelfId, &SV) -> Option<TV::SelfId> + Send + Sync + 'static,
    {
        self.decls.push(Decl::Fk {
            from: SV::NAME,
            to: TV::NAME,
            label: label.0,
            on_missing,
            build: Box::new(move |k, v| {
                let (Some(k), Some(v)) = (k.downcast_ref::<SV::SelfId>(), v.downcast_ref::<SV>())
                else {
                    return FkBuild::TypeMismatch;
                };
                match build(k, v) {
                    Some(tk) => FkBuild::Key(Box::new(tk)),
                    None => FkBuild::NoRef,
                }
            }),
        });
    }

    /// Denormalization: key projection + consistency predicate.
    pub fn denormalized<SV, DV, PF, EQ>(&mut self, project: PF, consistent: EQ)
    where
        SV: CollectionMeta + 'static,
        SV::SelfId: Key + 'static,
        DV: CollectionMeta + 'static,
        DV::SelfId: Key + 'static,
        PF: Fn(&SV::SelfId, &SV) -> DV::SelfId + Send + Sync + 'static,
        EQ: Fn(&SV, &DV) -> bool + Send + Sync + 'static,
    {
        self.decls.push(Decl::Denorm {
            source: SV::NAME,
            denorm: DV::NAME,
            project: Box::new(move |k, v| {
                let k = k.downcast_ref::<SV::SelfId>()?;
                let v = v.downcast_ref::<SV>()?;
                let dk = project(k, v);
                Some((dk.as_bytes().to_vec(), Box::new(dk) as Box<dyn Any>))
            }),
            consistent: Box::new(move |s, d| {
                Some(consistent(s.downcast_ref::<SV>()?, d.downcast_ref::<DV>()?))
            }),
        });
    }

    /// Counter: `field` in CV is recomputed as count of `source` records whose key
    /// starts with `prefix(cv key)`.
    pub fn counter<CV, RF, PF>(
        &mut self,
        field: &'static str,
        read: RF,
        source: &'static str,
        prefix: PF,
    ) where
        CV: CollectionMeta + 'static,
        CV::SelfId: Key + 'static,
        RF: Fn(&CV) -> u64 + Send + Sync + 'static,
        PF: Fn(&CV::SelfId) -> Vec<u8> + Send + Sync + 'static,
    {
        self.decls.push(Decl::Counter {
            collection: CV::NAME,
            field,
            source,
            read: Box::new(move |v| Some(read(v.downcast_ref::<CV>()?))),
            prefix: Box::new(move |k| Some(prefix(k.downcast_ref::<CV::SelfId>()?))),
        });
    }
}

impl BuiltRegistry {
    pub fn graph(&self) -> SchemaGraph {
        let mut g = SchemaGraph::default();
        for e in &self.reg.entries {
            g.collections.push(CollectionNode {
                name: e.name.to_string(),
                kind: e.kind,
                ty: e.typ,
                self_brand: e.self_brand.map(str::to_string),
                storage: Some(e.storage),
            });
        }
        for (name, kind) in &self.reg.external {
            g.collections.push(CollectionNode {
                name: name.to_string(),
                kind: *kind,
                ty: Typ::Custom("external", &[]),
                self_brand: None,
                storage: None,
            });
        }
        for e in &self.reg.entries {
            for (loc, r) in &e.auto_refs {
                if let Some(target) = self.brand_targets.get(r.brand) {
                    if e.kind == CollectionKind::Entity
                        && *loc == RefIn::Key
                        && e.self_brand == Some(r.brand)
                    {
                        continue;
                    }
                    g.relations.push(RelationEdge {
                        from: e.name.to_string(),
                        to: target.to_string(),
                        kind: RelationKind::Fk {
                            field_path: r.path.clone(),
                            brand: r.brand.to_string(),
                            on_missing: OnMissing::Error,
                            validated: true,
                        },
                    });
                }
            }
        }
        for d in &self.reg.decls {
            match d {
                Decl::Index { source, index, .. } => g.relations.push(RelationEdge {
                    from: source.to_string(),
                    to: index.to_string(),
                    kind: RelationKind::Index,
                }),
                Decl::Fk {
                    from,
                    to,
                    label,
                    on_missing,
                    ..
                } => g.relations.push(RelationEdge {
                    from: from.to_string(),
                    to: to.to_string(),
                    kind: RelationKind::Fk {
                        field_path: vec![],
                        brand: label.clone(),
                        on_missing: *on_missing,
                        validated: true,
                    },
                }),
                Decl::Denorm { source, denorm, .. } => g.relations.push(RelationEdge {
                    from: source.to_string(),
                    to: denorm.to_string(),
                    kind: RelationKind::Denorm,
                }),
                Decl::Counter {
                    collection,
                    field,
                    source,
                    ..
                } => g.relations.push(RelationEdge {
                    from: collection.to_string(),
                    to: source.to_string(),
                    kind: RelationKind::Counter {
                        field: field.to_string(),
                    },
                }),
            }
        }
        g
    }
}

fn find<'a>(entries: &'a [CollectionEntry], name: &str) -> Option<&'a CollectionEntry> {
    entries.iter().find(|e| e.name == name)
}

pub(super) fn run_decls(reg: &BuiltRegistry, report: &mut super::validate::ValidationReport) {
    use super::validate::{Finding, hex};
    // Expected index keys are accumulated across ALL Index decls targeting the
    // same index collection — several sources may legitimately feed one index.
    // The orphan pass runs once per index, after every source contributed.
    let mut expected_by_index: std::collections::HashMap<
        &'static str,
        std::collections::HashSet<Vec<u8>>,
    > = Default::default();
    for d in &reg.reg.decls {
        match d {
            Decl::Index {
                source,
                index,
                project,
                value_eq,
            } => {
                let Some(src) = reg.reg.entries.iter().find(|e| e.name == *source) else {
                    report.findings.push(Finding::RegistryError {
                        message: format!("index decl: source `{source}` not registered"),
                    });
                    continue;
                };
                let Some(idx) = reg.reg.entries.iter().find(|e| e.name == *index) else {
                    report.findings.push(Finding::RegistryError {
                        message: format!("index decl: index `{index}` not registered"),
                    });
                    continue;
                };

                // pass 1: source -> expected index entries
                let expected = expected_by_index.entry(*index).or_default();
                (src.scan_any)(&mut |_kb, k, v| {
                    let Some((ik_bytes, ik, iv)) = project(k, v) else {
                        report.findings.push(Finding::RegistryError {
                            message: format!(
                                "index decl `{source}`->`{index}`: type mismatch in projection"
                            ),
                        });
                        return;
                    };
                    expected.insert(ik_bytes.clone());
                    match (idx.with_value_any)(&*ik, &mut |actual| {
                        if value_eq(&*iv, actual) == Some(false) {
                            report.findings.push(Finding::IndexValueMismatch {
                                index: index.to_string(),
                                key: hex(&ik_bytes),
                            });
                        }
                    }) {
                        Some(true) => {}
                        Some(false) => report.findings.push(Finding::IndexMissing {
                            source: source.to_string(),
                            index: index.to_string(),
                            key: hex(&ik_bytes),
                        }),
                        None => report.findings.push(Finding::RegistryError {
                            message: format!("index decl `{index}`: key type mismatch"),
                        }),
                    }
                });
            }
            Decl::Fk {
                from,
                to,
                label,
                on_missing,
                build,
            } => {
                let (Some(src), Some(target)) =
                    (find(&reg.reg.entries, from), find(&reg.reg.entries, to))
                else {
                    report.findings.push(Finding::RegistryError {
                        message: format!("fk decl `{from}`->`{to}`: collection not registered"),
                    });
                    continue;
                };
                (src.scan_any)(&mut |kb, k, v| match build(k, v) {
                    FkBuild::NoRef => {}
                    FkBuild::TypeMismatch => report.findings.push(Finding::RegistryError {
                        message: format!("fk decl `{from}`->`{to}` ({label}): type mismatch"),
                    }),
                    FkBuild::Key(tk) => match (target.contains_any)(&*tk) {
                        Some(true) => {}
                        Some(false) => {
                            if *on_missing != OnMissing::AllowMissing {
                                report.findings.push(Finding::DanglingFk {
                                    collection: from.to_string(),
                                    key: hex(kb),
                                    target: to.to_string(),
                                    brand: label.clone(),
                                    id: 0,
                                    on_missing: *on_missing,
                                });
                            }
                        }
                        None => report.findings.push(Finding::RegistryError {
                            message: format!("fk decl `{to}`: key type mismatch"),
                        }),
                    },
                });
            }
            Decl::Denorm {
                source,
                denorm,
                project,
                consistent,
            } => {
                let (Some(src), Some(dn)) = (
                    find(&reg.reg.entries, source),
                    find(&reg.reg.entries, denorm),
                ) else {
                    report.findings.push(Finding::RegistryError {
                        message: format!(
                            "denorm decl `{source}`->`{denorm}`: collection not registered"
                        ),
                    });
                    continue;
                };
                (src.scan_any)(&mut |_kb, k, v| {
                    let Some((dk_bytes, dk)) = project(k, v) else {
                        report.findings.push(Finding::RegistryError {
                            message: format!(
                                "denorm decl `{source}`->`{denorm}`: type mismatch in projection"
                            ),
                        });
                        return;
                    };
                    match (dn.with_value_any)(&*dk, &mut |dv| {
                        if consistent(v, dv) == Some(false) {
                            report.findings.push(Finding::DenormMismatch {
                                source: source.to_string(),
                                denorm: denorm.to_string(),
                                key: hex(&dk_bytes),
                            });
                        }
                    }) {
                        Some(true) => {}
                        Some(false) => report.findings.push(Finding::DenormMissing {
                            source: source.to_string(),
                            denorm: denorm.to_string(),
                            key: hex(&dk_bytes),
                        }),
                        None => report.findings.push(Finding::RegistryError {
                            message: format!("denorm decl `{denorm}`: key type mismatch"),
                        }),
                    }
                });
            }
            Decl::Counter {
                collection,
                field,
                source,
                read,
                prefix,
            } => {
                let (Some(cv), Some(src)) = (
                    find(&reg.reg.entries, collection),
                    find(&reg.reg.entries, source),
                ) else {
                    report.findings.push(Finding::RegistryError {
                        message: format!(
                            "counter decl `{collection}`<-`{source}`: collection not registered"
                        ),
                    });
                    continue;
                };
                let mut counters: Vec<(Vec<u8>, u64, String)> = Vec::new();
                (cv.scan_any)(&mut |kb, k, v| {
                    if let (Some(p), Some(stored)) = (prefix(k), read(v)) {
                        counters.push((p, stored, hex(kb)));
                    }
                });
                let mut actual = vec![0u64; counters.len()];
                (src.scan_any)(&mut |kb, _k, _v| {
                    for (i, (p, _, _)) in counters.iter().enumerate() {
                        if kb.starts_with(p) {
                            actual[i] += 1;
                        }
                    }
                });
                for ((_, stored, key), got) in counters.into_iter().zip(actual) {
                    if stored != got {
                        report.findings.push(Finding::CounterMismatch {
                            collection: collection.to_string(),
                            key,
                            field: field.to_string(),
                            stored,
                            expected: got,
                        });
                    }
                }
            }
        }
    }

    // pass 2: orphan index entries — once per index collection, against the
    // union of expected keys from all its sources.
    for (index_name, expected) in &expected_by_index {
        let Some(idx) = reg.reg.entries.iter().find(|e| e.name == *index_name) else {
            continue; // RegistryError already reported in pass 1
        };
        (idx.scan_any)(&mut |kb, _k, _v| {
            if !expected.contains(kb) {
                report.findings.push(Finding::IndexOrphan {
                    index: index_name.to_string(),
                    key: hex(kb),
                });
            }
        });
    }
}