//! A semantic containment index shared by the sidebar and focused model browser.
//! Groups organize roots only; allocations never become containment. Traversals
//! are iterative and search keeps the ancestry of each result visible.
use oa_core::{Id, Kind, Layer, Project};
use std::collections::{BTreeMap, BTreeSet};

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Group {
    Participants,
    Work,
    MissionsCapabilities,
    Behavior,
    DataInterfaces,
    Other,
}

impl Group {
    pub const ALL: [Self; 6] = [
        Self::Participants,
        Self::Work,
        Self::MissionsCapabilities,
        Self::Behavior,
        Self::DataInterfaces,
        Self::Other,
    ];

    pub fn label(self, layer: Layer) -> &'static str {
        match (self, layer) {
            (Self::Participants, Layer::Oa) => "Entities, actors & roles",
            (Self::Participants, Layer::Sa) => "System & actors",
            (Self::Work, Layer::Oa) => "Operational activities",
            (Self::Work, Layer::Sa) => "System functions",
            (Self::MissionsCapabilities, _) => "Missions & capabilities",
            (Self::Behavior, _) => "Behavior & scenarios",
            (Self::DataInterfaces, _) => "Data & interfaces",
            (Self::Other, _) => "Other elements",
        }
    }

    fn for_root(kind: Kind) -> Self {
        match kind {
            Kind::Entity | Kind::Actor | Kind::Role | Kind::System | Kind::SystemActor => {
                Self::Participants
            }
            Kind::Activity | Kind::Function => Self::Work,
            Kind::Capability
            | Kind::OperationalMission
            | Kind::SystemCapability
            | Kind::Mission => Self::MissionsCapabilities,
            Kind::Process
            | Kind::FunctionalChain
            | Kind::Scenario
            | Kind::Lifeline
            | Kind::Occurrence
            | Kind::StateMachine
            | Kind::Region
            | Kind::Event => Self::Behavior,
            Kind::ExchangeItem
            | Kind::DataType
            | Kind::DataField
            | Kind::Interface
            | Kind::Port => Self::DataInterfaces,
            kind if kind.is_state() => Self::Behavior,
            _ => Self::Other,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Row {
    pub id: Id,
    /// Actual containment depth, with no display-specific truncation.
    pub depth: usize,
    pub has_children: bool,
    /// False only for an ancestor retained as context for a nonempty search.
    pub matched: bool,
}

struct SearchText {
    name: String,
    alias: String,
    kind: String,
}

pub struct Index {
    /// Every group is present, even when it has no roots in this perspective.
    pub roots: BTreeMap<Group, Vec<Id>>,
    children: BTreeMap<Id, Vec<Id>>,
    parents: BTreeMap<Id, Id>,
    groups: BTreeMap<Id, Group>,
    order: BTreeMap<Group, Vec<Row>>,
    search: BTreeMap<Id, SearchText>,
}

impl Index {
    pub fn new(project: &Project, layer: Layer) -> Self {
        let elements: BTreeMap<_, _> = project
            .elements
            .values()
            .filter(|element| element.layer() == layer)
            .map(|element| (element.id, element))
            .collect();
        let sort_keys: BTreeMap<_, _> = elements
            .iter()
            .map(|(&id, element)| (id, (element.name.to_lowercase(), element.alias.clone(), id)))
            .collect();
        let mut index = Self {
            roots: Group::ALL
                .into_iter()
                .map(|group| (group, vec![]))
                .collect(),
            children: BTreeMap::new(),
            parents: BTreeMap::new(),
            groups: BTreeMap::new(),
            order: BTreeMap::new(),
            search: elements
                .iter()
                .map(|(&id, element)| {
                    (
                        id,
                        SearchText {
                            name: element.name.to_lowercase(),
                            alias: element.alias.to_lowercase(),
                            kind: element.kind().label().to_lowercase(),
                        },
                    )
                })
                .collect(),
        };
        for (&id, element) in &elements {
            if let Some(parent) = element
                .parent
                .filter(|parent| elements.contains_key(parent))
            {
                index.parents.insert(id, parent);
                index.children.entry(parent).or_default().push(id);
            } else {
                index
                    .roots
                    .get_mut(&Group::for_root(element.kind()))
                    .unwrap()
                    .push(id);
            }
        }
        for siblings in index.roots.values_mut().chain(index.children.values_mut()) {
            siblings.sort_unstable_by(|left, right| sort_keys[left].cmp(&sort_keys[right]));
        }
        for group in Group::ALL {
            let mut stack: Vec<_> = index.roots[&group]
                .iter()
                .rev()
                .map(|&id| (id, 0usize))
                .collect();
            let mut order = vec![];
            while let Some((id, depth)) = stack.pop() {
                // Accepted models are acyclic; also bound traversal of a damaged
                // candidate so navigation itself cannot loop indefinitely.
                if index.groups.insert(id, group).is_some() {
                    continue;
                }
                let children = index.children(id);
                order.push(Row {
                    id,
                    depth,
                    has_children: !children.is_empty(),
                    matched: true,
                });
                stack.extend(children.iter().rev().map(|&child| (child, depth + 1)));
            }
            index.order.insert(group, order);
        }
        index
    }

    pub fn children(&self, id: Id) -> &[Id] {
        self.children.get(&id).map_or(&[], Vec::as_slice)
    }

    /// Breadcrumb identities in root-to-selection order, including selection.
    pub fn ancestry(&self, id: Id) -> Vec<Id> {
        if !self.search.contains_key(&id) {
            return vec![];
        }
        let mut result = vec![];
        let mut seen = BTreeSet::new();
        let mut current = Some(id);
        while let Some(id) = current {
            if !seen.insert(id) {
                break;
            }
            result.push(id);
            current = self.parents.get(&id).copied();
        }
        result.reverse();
        result
    }

    pub fn group(&self, id: Id) -> Option<Group> {
        self.groups.get(&id).copied()
    }

    /// Case-insensitive terms match names, aliases or kind labels. An `@term`
    /// matches aliases only. Search reveals matching paths without changing the
    /// user's collapsed state; unrelated descendants remain filtered out.
    pub fn rows(&self, group: Group, query: &str, collapsed_ids: &BTreeSet<Id>) -> Vec<Row> {
        let normalized = query.to_lowercase();
        let terms: Vec<_> = normalized.split_whitespace().collect();
        let order = &self.order[&group];
        if terms.is_empty() {
            let mut rows = vec![];
            let mut hidden_below = None;
            for &row in order {
                if hidden_below.is_some_and(|depth| row.depth > depth) {
                    continue;
                }
                hidden_below = collapsed_ids.contains(&row.id).then_some(row.depth);
                rows.push(row);
            }
            return rows;
        }
        let matched: BTreeSet<_> = order
            .iter()
            .filter(|row| {
                let text = &self.search[&row.id];
                terms.iter().all(|term| {
                    if let Some(alias) = term.strip_prefix('@') {
                        text.alias.contains(alias)
                    } else {
                        text.name.contains(term)
                            || text.alias.contains(term)
                            || text.kind.contains(term)
                    }
                })
            })
            .map(|row| row.id)
            .collect();
        let mut visible = matched.clone();
        // Reverse preorder propagates context to all ancestors in one pass;
        // walking every result's full path would be quadratic on deep models.
        for row in order.iter().rev() {
            if visible.contains(&row.id)
                && let Some(parent) = self.parents.get(&row.id)
            {
                visible.insert(*parent);
            }
        }
        order
            .iter()
            .filter(|row| visible.contains(&row.id))
            .map(|row| Row {
                matched: matched.contains(&row.id),
                ..*row
            })
            .collect()
    }
}

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

    fn add(project: &mut Project, kind: Kind, name: &str, parent: Option<Id>) -> Id {
        if kind == Kind::Port {
            let id = Id::new_v4();
            let port = oa_core::make_port(
                project,
                id,
                parent.unwrap(),
                name,
                oa_core::PortContract::Component {
                    kind: oa_core::ComponentPortKind::Flow,
                    orientation: oa_core::ComponentPortOrientation::Unset,
                },
            )
            .unwrap();
            project.elements.insert(id, port);
            return id;
        }
        project.add_element(kind, Layer::Sa, name, parent).unwrap()
    }

    fn ids(rows: Vec<Row>) -> Vec<Id> {
        rows.into_iter().map(|row| row.id).collect()
    }

    #[test]
    fn prefix_names_cannot_separate_a_parent_from_its_children() {
        let mut project = Project::empty();
        let parent = add(&mut project, Kind::SystemActor, "Rescue", None);
        let child = add(&mut project, Kind::Port, "Output", Some(parent));
        let other = add(&mut project, Kind::SystemActor, "Rescue archive", None);
        let index = Index::new(&project, Layer::Sa);
        assert_eq!(
            ids(index.rows(Group::Participants, "", &BTreeSet::new())),
            vec![parent, child, other]
        );
        assert_eq!(index.children(parent), &[child]);
        assert_eq!(index.ancestry(child), vec![parent, child]);
        assert_eq!(index.group(child), Some(Group::Participants));
        assert!(index.roots[&Group::DataInterfaces].is_empty());
    }

    #[test]
    fn deep_containment_and_search_keep_the_full_path_without_recursion() {
        let mut project = Project::empty();
        let mut parent = None;
        let mut chain = vec![];
        for depth in 0..260 {
            let id = add(
                &mut project,
                Kind::Function,
                &format!("Function {depth}"),
                parent,
            );
            parent = Some(id);
            chain.push(id);
        }
        let index = Index::new(&project, Layer::Sa);
        let rows = index.rows(Group::Work, "", &BTreeSet::new());
        assert_eq!(rows.len(), chain.len());
        assert_eq!(rows.last().unwrap().depth, 259);
        assert_eq!(index.ancestry(*chain.last().unwrap()), chain);
        let filtered = index.rows(Group::Work, "Function 259", &BTreeSet::from([chain[0]]));
        assert_eq!(filtered.len(), chain.len());
        assert!(filtered.last().unwrap().matched);
        assert!(filtered[..259].iter().all(|row| !row.matched));
    }

    #[test]
    fn collapse_and_search_retain_context_without_changing_collapse_state() {
        let mut project = Project::empty();
        let owner = add(&mut project, Kind::SystemActor, "Operator", None);
        let port = add(&mut project, Kind::Port, "Telemetry", Some(owner));
        project.elements.get_mut(&port).unwrap().alias = "report_out".into();
        let unrelated = add(&mut project, Kind::Port, "Power", Some(owner));
        let index = Index::new(&project, Layer::Sa);
        let collapsed = BTreeSet::from([owner]);
        assert_eq!(
            ids(index.rows(Group::Participants, "", &collapsed)),
            vec![owner]
        );
        let filtered = index.rows(Group::Participants, "@REPORT port", &collapsed);
        assert_eq!(ids(filtered.clone()), vec![owner, port]);
        assert!(!filtered[0].matched);
        assert!(filtered[1].matched);
        assert!(filtered[0].has_children);
        assert!(!filtered[1].has_children);
        assert!(!filtered.iter().any(|row| row.id == unrelated));
        assert_eq!(
            ids(index.rows(Group::Participants, "", &collapsed)),
            vec![owner]
        );
        assert!(
            index
                .rows(Group::Participants, "missing", &collapsed)
                .is_empty()
        );
    }

    #[test]
    fn equal_names_have_stable_alias_and_identity_tie_breakers() {
        let mut project = Project::empty();
        for (value, alias) in [(3, "zulu"), (2, "alpha"), (1, "alpha")] {
            let mut element = project
                .make_element(Kind::Function, Layer::Sa, "Same", None)
                .unwrap();
            element.id = Id::from_u128(value);
            element.alias = alias.into();
            project.elements.insert(element.id, element);
        }
        let index = Index::new(&project, Layer::Sa);
        assert_eq!(
            index.roots[&Group::Work],
            [Id::from_u128(1), Id::from_u128(2), Id::from_u128(3)]
        );
        assert_eq!(index.roots, Index::new(&project, Layer::Sa).roots);
    }

    #[test]
    fn perspective_isolation_and_allocations_do_not_change_containment() {
        let mut project = Project::empty();
        let actor = add(&mut project, Kind::SystemActor, "Operator", None);
        let work = add(&mut project, Kind::Function, "Respond", None);
        project.relate(oa_core::RelationKind::Allocation, "Performs", actor, work);
        let oa = project
            .add_element(Kind::Activity, Layer::Oa, "Respond", None)
            .unwrap();
        let oa_index = Index::new(&project, Layer::Oa);
        let sa_index = Index::new(&project, Layer::Sa);
        assert_eq!(oa_index.roots[&Group::Work], vec![oa]);
        assert_eq!(sa_index.roots[&Group::Work], vec![work]);
        assert!(sa_index.children(actor).is_empty());
        assert_eq!(sa_index.ancestry(work), vec![work]);
        assert_eq!(oa_index.group(work), None);
        assert_eq!(sa_index.group(oa), None);
        assert!(oa_index.ancestry(work).is_empty());
        assert_eq!(Group::Work.label(Layer::Oa), "Operational activities");
        assert_eq!(Group::Work.label(Layer::Sa), "System functions");
    }
}