armdb 0.14.2

//! Static FK assembly plan per collection (built in `finish`, consumed in `validate`).

use std::collections::{HashMap, HashSet};

use armour_core::{CompositeRef, FieldRef, PathSeg, RowPart};

use super::registry::Home;

pub type IdMap = HashMap<(Vec<PathSeg>, &'static str, Vec<PathSeg>), u64>;

/// Which planned edges to evaluate for this id snapshot.
#[derive(Debug, Clone)]
pub enum IdFlush {
    /// A `PathSeg::Elem` scope instance just completed — only edges with this exact scope.
    Scope(Vec<PathSeg>),
    /// End of row — edges whose scope is not an element instance (no trailing/internal per-elem).
    Final,
}

impl IdFlush {
    pub fn matches_edge_scope(&self, edge_scope: &[PathSeg]) -> bool {
        match self {
            // `scope` is the visitor stack at `leave(Elem)` (ends with `Elem`).
            // Prefix, not equality: scopes *inside* the element (an enum variant
            // or a nested struct under `Vec<T>`) are only ever complete at this
            // flush too, and the ids for this instance are dropped right after.
            IdFlush::Scope(scope) => edge_scope.starts_with(scope),
            // Skip edges whose scope is a container-element instance (contains `Elem`).
            IdFlush::Final => !edge_scope.iter().any(|s| matches!(s, PathSeg::Elem)),
        }
    }
}

pub type VisitRefsFn = Box<dyn Fn(&mut dyn FnMut(&[u8], &IdMap, &IdMap, IdFlush)) + Send + Sync>;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AmbiguityReason {
    RepeatedBrand {
        brand: String,
        paths: Vec<Vec<PathSeg>>,
    },
    RepeatedBrandHomes {
        homes: Vec<String>,
    },
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AmbiguousRef {
    pub collection: String,
    pub part: RowPart,
    pub scope: Vec<PathSeg>,
    /// The target shape as a **sorted** brand multiset — this finding is about a
    /// shape that could not be matched, so no home's field order applies.
    /// (`RefDesc::Assembled.shape`, which describes an assembled reference, is
    /// in the home's field order instead.)
    pub shape: Vec<String>,
    pub reason: AmbiguityReason,
}

#[derive(Clone)]
pub(in crate::schema) enum IdSlot {
    AtPath {
        part: RowPart,
        path: Vec<PathSeg>,
        brand: &'static str,
    },
    KeyRoot {
        brand: &'static str,
    },
}

#[derive(Clone)]
pub(in crate::schema) struct PlannedEdge {
    pub part: RowPart,
    /// Scope this edge lives in — matched against the flush that just completed.
    pub scope: Vec<PathSeg>,
    /// Plan-time reference shape (field order of the first home used to build slots).
    pub home_shape: Vec<&'static str>,
    /// Parallel to `home_shapes`: home collection names.
    pub home_names: Vec<&'static str>,
    /// Each home's own `KeyShape::SHAPE` (field order), parallel to `home_names`.
    pub home_shapes: Vec<Vec<&'static str>>,
    pub slots: Vec<IdSlot>,
    pub reference_path: Vec<PathSeg>,
}

#[derive(Clone)]
pub(in crate::schema) struct MatchPlan {
    pub ambiguities: Vec<AmbiguousRef>,
    pub edges: Vec<PlannedEdge>,
    pub unvalidated_brands: HashSet<&'static str>,
}

pub(in crate::schema) fn scope_of_path(path: &[PathSeg]) -> Vec<PathSeg> {
    // `Elem` is a per-instance scope opener. When it is also the leaf
    // (`Vec<Fuid>`), keep it so plan scopes match the runtime visitor stack
    // flushed on `leave(Elem)`.
    if path.last() == Some(&PathSeg::Elem) {
        return path.to_vec();
    }
    path[..path.len().saturating_sub(1)].to_vec()
}

fn scope_of_composite(c: &CompositeRef) -> Vec<PathSeg> {
    if c.parts.is_empty() {
        return Vec::new();
    }
    let mut common = scope_of_path(&c.parts[0]);
    for part in c.parts.iter().skip(1) {
        let scope = scope_of_path(part);
        let mut i = 0;
        while i < common.len() && i < scope.len() && common[i] == scope[i] {
            i += 1;
        }
        common.truncate(i);
    }
    common
}

fn sorted_shape_key(shape: &[&'static str]) -> Box<[&'static str]> {
    let mut v: Vec<_> = shape.to_vec();
    v.sort_unstable();
    v.into_boxed_slice()
}

/// How many ids of each brand a shape needs (a shape may repeat a brand).
fn brand_multiplicities(shape: &[&'static str]) -> HashMap<&'static str, usize> {
    let mut m: HashMap<&'static str, usize> = HashMap::new();
    for b in shape {
        *m.entry(*b).or_default() += 1;
    }
    m
}

fn sorted_shape_vec(shape: &[&'static str]) -> Vec<String> {
    let mut v: Vec<String> = shape.iter().map(|s| (*s).to_string()).collect();
    v.sort();
    v
}

fn shapes_multiset_eq(a: &[&'static str], b: &[&'static str]) -> bool {
    sorted_shape_key(a) == sorted_shape_key(b)
}

fn brand_at_path(
    auto_refs: &[(RowPart, FieldRef)],
    part: RowPart,
    path: &[PathSeg],
) -> Option<&'static str> {
    auto_refs
        .iter()
        .find(|(p, r)| *p == part && r.path == path)
        .map(|(_, r)| r.brand)
}

fn brands_in_scope<'a>(
    refs: impl Iterator<Item = &'a FieldRef>,
) -> HashMap<&'static str, Vec<Vec<PathSeg>>> {
    let mut m: HashMap<&'static str, Vec<Vec<PathSeg>>> = HashMap::new();
    for r in refs {
        m.entry(r.brand).or_default().push(r.path.clone());
    }
    m
}

fn key_root_refs(auto_refs: &[(RowPart, FieldRef)]) -> Vec<&FieldRef> {
    auto_refs
        .iter()
        .filter(|(p, r)| *p == RowPart::Key && scope_of_path(&r.path).is_empty())
        .map(|(_, r)| r)
        .collect()
}

fn path_for_brand_in_key(refs: &[&FieldRef], brand: &str) -> Option<Vec<PathSeg>> {
    refs.iter()
        .find(|r| r.brand == brand)
        .map(|r| r.path.clone())
}

fn reference_path_for_shape(
    home_shape: &[&str],
    slots: &[IdSlot],
    key_roots: &[&FieldRef],
) -> Vec<PathSeg> {
    let last_brand = home_shape.last().copied().unwrap_or("");
    for slot in slots.iter().rev() {
        match slot {
            IdSlot::AtPath { path, brand, .. } if *brand == last_brand => return path.clone(),
            IdSlot::KeyRoot { brand } if *brand == last_brand => {
                return path_for_brand_in_key(key_roots, brand).unwrap_or_default();
            }
            _ => {}
        }
    }
    match slots.last() {
        Some(IdSlot::AtPath { path, .. }) => path.clone(),
        Some(IdSlot::KeyRoot { brand }) => {
            path_for_brand_in_key(key_roots, brand).unwrap_or_default()
        }
        None => Vec::new(),
    }
}

/// Graph export: last component path for a home's field order (brands unique in slots).
pub(in crate::schema) fn reference_path_for_graph(
    home_shape: &[&str],
    slots: &[IdSlot],
) -> Vec<PathSeg> {
    let last_brand = home_shape.last().copied().unwrap_or("");
    for slot in slots {
        match slot {
            IdSlot::AtPath { path, brand, .. } if *brand == last_brand => return path.clone(),
            IdSlot::KeyRoot { brand } if *brand == last_brand => return Vec::new(),
            _ => {}
        }
    }
    match slots.last() {
        Some(IdSlot::AtPath { path, .. }) => path.clone(),
        _ => Vec::new(),
    }
}

fn suppress_self_shape(
    part: RowPart,
    scope: &[PathSeg],
    self_shape: Option<&[&'static str]>,
    candidate_shape: &[&'static str],
) -> bool {
    part == RowPart::Key
        && scope.is_empty()
        && self_shape.is_some_and(|s| shapes_multiset_eq(s, candidate_shape))
}

fn home_matches_decl(
    part_brands: &[&'static str],
    with_key: bool,
    key_roots: &[&FieldRef],
    home_shape: &[&'static str],
) -> bool {
    let mut remaining: Vec<&'static str> = home_shape.to_vec();
    for b in part_brands {
        if let Some(i) = remaining.iter().position(|x| *x == *b) {
            remaining.remove(i);
        } else {
            return false;
        }
    }
    if remaining.is_empty() {
        return true;
    }
    if !with_key {
        return false;
    }
    // What the declaration did not name comes from the row-key root. The key
    // must offer at least as many ids per brand as the shape still needs;
    // offering *more* is an ambiguity (reported by the caller), not a mismatch —
    // deciding it here would make that finding unreachable.
    brand_multiplicities(&remaining)
        .into_iter()
        .all(|(b, n)| key_roots.iter().filter(|kr| kr.brand == b).count() >= n)
}

fn remaining_key_brands(
    part_brands: &[&'static str],
    home_shape: &[&'static str],
) -> Vec<&'static str> {
    let mut remaining: Vec<&'static str> = home_shape.to_vec();
    for b in part_brands {
        if let Some(i) = remaining.iter().position(|x| *x == *b) {
            remaining.remove(i);
        }
    }
    remaining
}

fn slot_for_brand(
    part: RowPart,
    brand: &'static str,
    refs_in_scope: &[&FieldRef],
    decl_part_paths: &[(Vec<PathSeg>, &'static str)],
) -> IdSlot {
    if let Some((path, _)) = decl_part_paths.iter().find(|(_, b)| *b == brand) {
        return IdSlot::AtPath {
            part,
            path: path.clone(),
            brand,
        };
    }
    if let Some(r) = refs_in_scope.iter().find(|r| r.brand == brand) {
        return IdSlot::AtPath {
            part,
            path: r.path.clone(),
            brand,
        };
    }
    IdSlot::KeyRoot { brand }
}

fn key_paths_for_brand(key_roots: &[&FieldRef], brand: &str) -> Vec<Vec<PathSeg>> {
    key_roots
        .iter()
        .filter(|kr| kr.brand == brand)
        .map(|kr| kr.path.clone())
        .collect()
}

type AmbiguityKey = (RowPart, Vec<PathSeg>, Vec<String>);

fn push_ambiguity_unique(
    ambiguities: &mut Vec<AmbiguousRef>,
    seen: &mut HashSet<AmbiguityKey>,
    amb: AmbiguousRef,
) {
    let key = (amb.part, amb.scope.clone(), amb.shape.clone());
    if seen.insert(key) {
        ambiguities.push(amb);
    }
}

fn mark_shape_validated(validated: &mut HashSet<&'static str>, shape: &[&'static str]) {
    for brand in shape {
        if !brand.is_empty() {
            validated.insert(*brand);
        }
    }
}

fn record_edge(
    edges: &mut Vec<PlannedEdge>,
    validated: &mut HashSet<&'static str>,
    part: RowPart,
    scope: Vec<PathSeg>,
    home_list: &[&Home],
    slots: Vec<IdSlot>,
    key_roots: &[&FieldRef],
) {
    let home = &home_list[0];
    for slot in &slots {
        match slot {
            IdSlot::AtPath { brand, .. } | IdSlot::KeyRoot { brand } => {
                validated.insert(*brand);
            }
        }
    }
    let reference_path = reference_path_for_shape(home.shape, &slots, key_roots);
    edges.push(PlannedEdge {
        part,
        scope,
        home_shape: home.shape.to_vec(),
        home_names: home_list.iter().map(|h| h.name).collect(),
        home_shapes: home_list.iter().map(|h| h.shape.to_vec()).collect(),
        slots,
        reference_path,
    });
}

fn build_declaration_edge(
    part: RowPart,
    decl: &CompositeRef,
    auto_refs: &[(RowPart, FieldRef)],
    refs_in_scope: &[&FieldRef],
    home_list: &[&Home],
    has_repeat: bool,
) -> Option<Vec<IdSlot>> {
    let home = &home_list[0];
    if home.shape.iter().any(|b| b.is_empty()) {
        return None;
    }

    let mut decl_entries: Vec<(Vec<PathSeg>, &'static str)> = Vec::new();
    for part_path in &decl.parts {
        let brand = brand_at_path(auto_refs, part, part_path)?;
        decl_entries.push((part_path.clone(), brand));
    }

    if has_repeat {
        // Positional: declaration parts in order onto SHAPE positions.
        let mut slots = Vec::with_capacity(home.shape.len());
        for (i, brand_in_shape) in home.shape.iter().enumerate() {
            if i < decl.parts.len() {
                let path = &decl.parts[i];
                slots.push(IdSlot::AtPath {
                    part,
                    path: path.clone(),
                    brand: brand_in_shape,
                });
            } else {
                slots.push(IdSlot::KeyRoot {
                    brand: brand_in_shape,
                });
            }
        }
        return Some(slots);
    }

    let mut slots = Vec::with_capacity(home.shape.len());
    for brand_in_shape in home.shape {
        slots.push(slot_for_brand(
            part,
            brand_in_shape,
            refs_in_scope,
            &decl_entries,
        ));
    }
    Some(slots)
}

pub(in crate::schema) fn build_match_plan(
    collection: &'static str,
    self_shape: Option<&[&'static str]>,
    auto_refs: &[(RowPart, FieldRef)],
    composites: &[(RowPart, CompositeRef)],
    homes: &HashMap<Box<[&'static str]>, Vec<Home>>,
) -> MatchPlan {
    let mut ambiguities = Vec::new();
    let mut edges = Vec::new();
    let mut validated_brands: HashSet<&'static str> = HashSet::new();
    let mut all_brands: HashSet<&'static str> = HashSet::new();
    let key_roots = key_root_refs(auto_refs);

    let mut amb_seen: HashSet<AmbiguityKey> = HashSet::new();
    let mut scopes: HashSet<(RowPart, Vec<PathSeg>)> = HashSet::new();
    for (part, r) in auto_refs {
        all_brands.insert(r.brand);
        scopes.insert((*part, scope_of_path(&r.path)));
    }
    for (part, c) in composites {
        scopes.insert((*part, scope_of_composite(c)));
    }

    let mut scope_list: Vec<(RowPart, Vec<PathSeg>)> = scopes.into_iter().collect();
    scope_list.sort();

    let mut home_entries: Vec<(&Box<[&'static str]>, &Vec<Home>)> = homes.iter().collect();
    home_entries.sort_by(|a, b| a.0.cmp(b.0));

    for (part, scope) in scope_list {
        let refs_in_scope: Vec<&FieldRef> = auto_refs
            .iter()
            .filter(|(p, r)| *p == part && scope_of_path(&r.path) == scope)
            .map(|(_, r)| r)
            .collect();

        let decls_in_scope: Vec<&CompositeRef> = composites
            .iter()
            .filter(|(p, c)| *p == part && scope_of_composite(c) == scope)
            .map(|(_, c)| c)
            .collect();

        let mut declared_shapes: HashSet<Vec<String>> = HashSet::new();

        for decl in &decls_in_scope {
            let part_brands: Vec<&'static str> = decl
                .parts
                .iter()
                .filter_map(|p| brand_at_path(auto_refs, part, p))
                .collect();
            if part_brands.is_empty() && !decl.with_key {
                continue;
            }

            // Repeats inside the declaration's own parts — not the whole scope.
            let mut decl_brand_counts: HashMap<&'static str, usize> = HashMap::new();
            for b in &part_brands {
                *decl_brand_counts.entry(*b).or_default() += 1;
            }
            let decl_has_repeat = decl_brand_counts.values().any(|&c| c > 1);

            for (_sorted_key, home_list) in &home_entries {
                if !home_matches_decl(&part_brands, decl.with_key, &key_roots, home_list[0].shape) {
                    continue;
                }
                let sorted = sorted_shape_vec(home_list[0].shape);
                // Multi-component declarations suppress auto for that shape (B1: len-1 never suppresses).
                if sorted.len() > 1 {
                    declared_shapes.insert(sorted.clone());
                }

                let home_refs: Vec<&Home> = home_list.iter().collect();
                if home_refs.is_empty() {
                    continue;
                }

                // A single-component shape is already one edge per id in the
                // automatic pass below (it is never suppressed) — planning it
                // here too would duplicate every finding and graph edge.
                if home_list[0].shape.len() == 1 {
                    continue;
                }

                // The collection *is* this home: its own key root is not a
                // reference to itself. Checked before any ambiguity is raised,
                // otherwise a home whose key repeats a brand accuses itself.
                if suppress_self_shape(part, &scope, self_shape, home_refs[0].shape) {
                    mark_shape_validated(&mut validated_brands, home_refs[0].shape);
                    continue;
                }

                let key_needed = remaining_key_brands(&part_brands, home_list[0].shape);

                if decl.with_key
                    && let Some(&brand) = key_needed
                        .iter()
                        .find(|&&b| key_roots.iter().filter(|kr| kr.brand == b).count() > 1)
                {
                    push_ambiguity_unique(
                        &mut ambiguities,
                        &mut amb_seen,
                        AmbiguousRef {
                            collection: collection.to_string(),
                            part,
                            scope: scope.clone(),
                            shape: sorted.clone(),
                            reason: AmbiguityReason::RepeatedBrand {
                                brand: brand.to_string(),
                                paths: key_paths_for_brand(&key_roots, brand),
                            },
                        },
                    );
                    continue;
                }

                if decl_has_repeat && home_list.len() >= 2 {
                    push_ambiguity_unique(
                        &mut ambiguities,
                        &mut amb_seen,
                        AmbiguousRef {
                            collection: collection.to_string(),
                            part,
                            scope: scope.clone(),
                            shape: sorted,
                            reason: AmbiguityReason::RepeatedBrandHomes {
                                homes: home_list.iter().map(|h| h.name.to_string()).collect(),
                            },
                        },
                    );
                    if !decl.with_key {
                        break;
                    }
                    continue;
                }

                if let Some(slots) = build_declaration_edge(
                    part,
                    decl,
                    auto_refs,
                    &refs_in_scope,
                    &home_refs,
                    decl_has_repeat,
                ) {
                    record_edge(
                        &mut edges,
                        &mut validated_brands,
                        part,
                        scope.clone(),
                        &home_refs,
                        slots,
                        &key_roots,
                    );
                }
                // A1: `ref_with_key` may complete several home shapes — plan every match.
                // Fixed-part `ref(...)` matches at most one multiset; stop after the first.
                if !decl.with_key {
                    break;
                }
            }
        }

        let brand_counts = brands_in_scope(refs_in_scope.iter().copied());

        for (_sorted_key, home_list) in &home_entries {
            let shape_brands: Vec<&str> = home_list[0].shape.to_vec();
            let sorted_str: Vec<String> = {
                let mut v: Vec<String> = shape_brands.iter().map(|s| (*s).to_string()).collect();
                v.sort();
                v
            };
            // B1: single-component shapes always auto-match (one edge per field).
            if shape_brands.len() > 1 && declared_shapes.contains(&sorted_str) {
                continue;
            }
            if home_list
                .iter()
                .any(|h| h.shape.iter().any(|b| b.is_empty()))
            {
                continue;
            }
            let home_refs: Vec<&Home> = home_list.iter().collect();

            if shape_brands.len() == 1 {
                let brand = shape_brands[0];
                for r in &refs_in_scope {
                    if r.brand != brand {
                        continue;
                    }
                    if suppress_self_shape(part, &scope, self_shape, &[brand]) {
                        mark_shape_validated(&mut validated_brands, &[brand]);
                        continue;
                    }
                    record_edge(
                        &mut edges,
                        &mut validated_brands,
                        part,
                        scope.clone(),
                        &home_refs,
                        vec![IdSlot::AtPath {
                            part,
                            path: r.path.clone(),
                            brand,
                        }],
                        &key_roots,
                    );
                }
                continue;
            }

            // Multiset, not presence: a shape that repeats a brand needs that
            // many ids in the scope. Fewer than needed is plan rule 3's "some
            // brand missing" — no edge, no finding. (Presence alone would let a
            // `[space, msg, msg]` home assemble from a single `msg`, filling two
            // slots from one field and reporting a consistent row as dangling.)
            let needed = brand_multiplicities(home_list[0].shape);
            if needed
                .iter()
                .any(|(b, n)| brand_counts.get(*b).map_or(0, |p| p.len()) < *n)
            {
                continue;
            }

            // Own key shape: the key root of a home is not a reference to
            // itself. Before the ambiguity check, or a home whose own key
            // repeats a brand would raise an error against itself.
            if suppress_self_shape(part, &scope, self_shape, home_refs[0].shape) {
                mark_shape_validated(&mut validated_brands, home_refs[0].shape);
                continue;
            }

            // More candidates than the shape has slots for that brand (which a
            // repeated brand always implies): automatic assembly cannot choose
            // — a `ref(...)` declaration must.
            if let Some(brand) = shape_brands
                .iter()
                .find(|b| brand_counts.get(**b).is_some_and(|p| p.len() > 1))
            {
                push_ambiguity_unique(
                    &mut ambiguities,
                    &mut amb_seen,
                    AmbiguousRef {
                        collection: collection.to_string(),
                        part,
                        scope: scope.clone(),
                        shape: sorted_str,
                        reason: AmbiguityReason::RepeatedBrand {
                            brand: (*brand).to_string(),
                            paths: brand_counts[brand].clone(),
                        },
                    },
                );
                continue;
            }

            let mut slots = Vec::with_capacity(home_refs[0].shape.len());
            for brand_in_shape in home_refs[0].shape {
                // Unreachable given the multiset check above; degrade to
                // "no edge" rather than panicking inside `finish()`.
                let Some(r) = refs_in_scope.iter().find(|r| r.brand == *brand_in_shape) else {
                    slots.clear();
                    break;
                };
                slots.push(IdSlot::AtPath {
                    part,
                    path: r.path.clone(),
                    brand: brand_in_shape,
                });
            }
            if slots.len() != home_refs[0].shape.len() {
                continue;
            }
            record_edge(
                &mut edges,
                &mut validated_brands,
                part,
                scope.clone(),
                &home_refs,
                slots,
                &key_roots,
            );
        }
    }

    let unvalidated_brands = all_brands
        .into_iter()
        .filter(|b| !validated_brands.contains(b))
        .collect();

    MatchPlan {
        ambiguities,
        edges,
        unvalidated_brands,
    }
}

pub(in crate::schema) fn resolve_slot(slot: &IdSlot, key: &IdMap, value: &IdMap) -> Option<u64> {
    match slot {
        IdSlot::AtPath { part, path, brand } => {
            let map = match part {
                RowPart::Key => key,
                RowPart::Value => value,
            };
            map.get(&(scope_of_path(path), *brand, path.clone()))
                .copied()
        }
        IdSlot::KeyRoot { brand } => key
            .iter()
            .find(|((scope, b, _), _)| scope.is_empty() && *b == *brand)
            .map(|(_, v)| *v),
    }
}

pub(in crate::schema) fn gather_ids(
    slots: &[IdSlot],
    key: &IdMap,
    value: &IdMap,
) -> Option<Vec<u64>> {
    let mut ids = Vec::with_capacity(slots.len());
    for slot in slots {
        ids.push(resolve_slot(slot, key, value)?);
    }
    Some(ids)
}

fn slot_brand(slot: &IdSlot) -> &'static str {
    match slot {
        IdSlot::AtPath { brand, .. } | IdSlot::KeyRoot { brand } => brand,
    }
}

/// Reorder assembled ids from the plan slot order into a specific home's `SHAPE` order.
///
/// When `home_shape` matches the plan's reference shape, returns `ids` unchanged.
/// When brands are unique, reorders by brand assignment. Repeated brands with
/// differing field order across homes are not reorderable (caller skips that home).
pub(in crate::schema) fn ids_for_home_shape(
    slots: &[IdSlot],
    ids: &[u64],
    home_shape: &[&'static str],
    plan_shape: &[&'static str],
) -> Option<Vec<u64>> {
    if home_shape == plan_shape {
        return Some(ids.to_vec());
    }

    let mut seen_brands = HashSet::new();
    for slot in slots {
        if !seen_brands.insert(slot_brand(slot)) {
            return None;
        }
    }

    debug_assert_eq!(slots.len(), ids.len());
    let mut by_brand: HashMap<&'static str, u64> = HashMap::with_capacity(slots.len());
    for (slot, &id) in slots.iter().zip(ids) {
        by_brand.insert(slot_brand(slot), id);
    }
    let mut reordered = Vec::with_capacity(home_shape.len());
    for &brand in home_shape {
        reordered.push(*by_brand.get(brand)?);
    }
    Some(reordered)
}