use crate::*;
use std::collections::BTreeSet;

/// Every live reference in a typed element, including structural ownership.
pub fn element_references(e: &Element) -> Vec<Id> {
    let mut copy = e.clone();
    let mut refs = Vec::new();
    visit_element_refs(&mut copy, |id| refs.push(*id));
    refs
}
pub fn relation_references(r: &Relation) -> Vec<Id> {
    let mut copy = r.clone();
    let mut refs = Vec::new();
    visit_relation_refs(&mut copy, |id| refs.push(*id));
    refs
}
pub fn visit_element_refs(e: &mut Element, mut visit: impl FnMut(&mut Id)) {
    if let Some(parent) = &mut e.parent {
        visit(parent);
    }
    match &mut e.body {
        ElementBody::Oa(OaBody::Process { steps })
        | ElementBody::Sa(SaBody::FunctionalChain { steps }) => {
            for id in steps {
                visit(id);
            }
        }
        _ => {}
    }
    if let Some(shared) = e.shared_mut() {
        match shared {
            SharedBody::DataField {
                type_ref: TypeRef::Named(id),
                ..
            } => visit(id),
            SharedBody::Interface { items } => {
                for id in items {
                    visit(id);
                }
            }
            SharedBody::Event {
                body: EventBody::Signal { items },
            } => {
                for id in items {
                    visit(id);
                }
            }
            SharedBody::Port { contract, legacy } => {
                match contract {
                    PortContract::FunctionInput { incoming_items } => {
                        incoming_items.iter_mut().for_each(&mut visit)
                    }
                    PortContract::FunctionOutput { outgoing_items } => {
                        outgoing_items.iter_mut().for_each(&mut visit)
                    }
                    _ => {}
                }
                if let Some(id) = legacy.as_mut().and_then(|old| old.interface.as_mut()) {
                    visit(id);
                }
            }
            SharedBody::Scenario { context, steps, .. } => {
                if let Some(id) = context {
                    visit(id);
                }
                for id in steps {
                    visit(id);
                }
            }
            SharedBody::Lifeline {
                represented: Some(id),
            } => visit(id),
            SharedBody::Occurrence { body, .. } => match body {
                OccurrenceBody::Exchange {
                    exchange,
                    sender,
                    receiver,
                    reply_to,
                    values,
                    ..
                } => {
                    for id in [exchange, sender, receiver, reply_to].into_iter().flatten() {
                        visit(id);
                    }
                    for value in values {
                        visit(&mut value.field);
                    }
                }
                OccurrenceBody::Action { action, lifeline } => {
                    for id in [action, lifeline].into_iter().flatten() {
                        visit(id);
                    }
                }
                OccurrenceBody::SequenceReference { sequence } => {
                    if let Some(id) = sequence {
                        visit(id);
                    }
                }
                OccurrenceBody::Control { .. } => {}
                OccurrenceBody::ScenarioUse { scenario, bindings } => {
                    if let Some(id) = scenario {
                        visit(id);
                    }
                    for binding in bindings {
                        visit(&mut binding.referenced);
                        visit(&mut binding.local);
                    }
                }
                OccurrenceBody::State { state, lifeline } => {
                    for id in [state, lifeline].into_iter().flatten() {
                        visit(id);
                    }
                }
                OccurrenceBody::Fragment { branches, .. } => {
                    for id in branches {
                        visit(id);
                    }
                }
                OccurrenceBody::Branch { occurrences, .. } => {
                    for id in occurrences {
                        visit(id);
                    }
                }
            },
            SharedBody::Constraint { constrained, .. } => {
                for id in constrained {
                    visit(id);
                }
            }
            _ => {}
        }
    }
}
pub fn visit_relation_refs(r: &mut Relation, mut visit: impl FnMut(&mut Id)) {
    let (mut source, mut target) = r.body.endpoints();
    visit(&mut source);
    visit(&mut target);
    r.body.set_endpoints(source, target);
    match &mut r.body {
        RelationBody::StateTransition {
            triggers, effects, ..
        } => {
            for id in triggers.iter_mut().chain(effects) {
                visit(id);
            }
        }
        RelationBody::Interaction { items, .. } | RelationBody::BoundaryExchange { items, .. } => {
            for id in items {
                visit(id);
            }
        }
        RelationBody::FunctionalExchange {
            source_port,
            target_port,
            items,
            ..
        } => {
            for id in [source_port, target_port] {
                visit(id);
            }
            for id in items {
                visit(id);
            }
        }
        _ => {}
    }
}
pub fn remap_references(p: &mut Project, old: Id, new: Id) {
    for e in p.elements.values_mut() {
        visit_element_refs(e, |id| {
            if *id == old {
                *id = new;
            }
        });
    }
    for r in p.relations.values_mut() {
        visit_relation_refs(r, |id| {
            if *id == old {
                *id = new;
            }
        });
    }
    for d in p.diagrams.values_mut() {
        for n in d.nodes.values_mut() {
            if n.element == old {
                n.element = new;
            }
        }
    }
    for a in p.artifact_links.values_mut() {
        if a.owner == old {
            a.owner = new;
        }
    }
    for t in p.traces.values_mut() {
        if t.oa == old {
            t.oa = new;
        }
        if t.sa == old {
            t.sa = new;
        }
    }
}
/// IDs removed by a cascade. Cross-layer traced counterparts are never included.
pub fn deletion_impact(p: &Project, id: Id) -> BTreeSet<Id> {
    deletion_impact_inner(p, id, false)
}
fn deletion_impact_inner(p: &Project, id: Id, defer_exchange_allocations: bool) -> BTreeSet<Id> {
    let mut removed = BTreeSet::from([id]);
    loop {
        let before = removed.len();
        for e in p.elements.values() {
            let dependent_type = matches!(e.shared(),Some(SharedBody::DataField{type_ref:TypeRef::Named(t),..}) if removed.contains(t));
            if e.parent.is_some_and(|parent| removed.contains(&parent)) || dependent_type {
                removed.insert(e.id);
            }
            if matches!(e.shared(),Some(SharedBody::Occurrence{body:OccurrenceBody::Exchange{exchange:Some(r),..},..}) if removed.contains(r))
            {
                removed.insert(e.id);
            }
        }
        for r in p.relations.values() {
            let port_removed = matches!(&r.body, RelationBody::FunctionalExchange { source_port, target_port, .. } if removed.contains(source_port) || removed.contains(target_port));
            let work_allocation_removed = r.kind() == RelationKind::PortAllocation
                && p.relations.values().any(|allocation| {
                    allocation.kind() == RelationKind::Allocation
                        && removed.contains(&allocation.id)
                        && p.elements.get(&r.source()).and_then(|e| e.parent)
                            == Some(allocation.source())
                        && p.elements.get(&r.target()).and_then(|e| e.parent)
                            == Some(allocation.target())
                });
            if removed.contains(&r.source())
                || removed.contains(&r.target())
                || port_removed
                || work_allocation_removed
            {
                removed.insert(r.id);
            }
        }
        // Exchange allocation depends on both explicit endpoint allocations.
        for r in p
            .relations
            .values()
            .filter(|r| r.kind() == RelationKind::ExchangeAllocation)
        {
            let (Some(boundary), Some(work)) =
                (p.relations.get(&r.source()), p.relations.get(&r.target()))
            else {
                continue;
            };
            let RelationBody::FunctionalExchange {
                source_port,
                target_port,
                ..
            } = &work.body
            else {
                continue;
            };
            if !defer_exchange_allocations
                && p.relations.values().any(|allocation| {
                    allocation.kind() == RelationKind::PortAllocation
                        && removed.contains(&allocation.id)
                        && [boundary.source(), boundary.target()].contains(&allocation.source())
                        && [*source_port, *target_port].contains(&allocation.target())
                })
                && crate::ports::missing_exchange_port_allocations(p, r, &removed)
            {
                removed.insert(r.id);
            }
        }
        for d in p.diagrams.values() {
            for n in d.nodes.values() {
                if removed.contains(&d.id) || removed.contains(&n.element) {
                    removed.insert(n.id);
                }
            }
        }
        for t in p.traces.values() {
            if removed.contains(&t.oa) || removed.contains(&t.sa) {
                removed.insert(t.id);
            }
        }
        for a in p.artifact_links.values() {
            if removed.contains(&a.owner) {
                removed.insert(a.id);
            }
        }
        if removed.len() == before {
            break;
        }
    }
    removed
}
pub fn delete_cascade(p: &mut Project, id: Id) {
    let removed = deletion_impact(p, id);
    remove_objects(p, removed);
}
/// Preserve conditional exchange-allocation dependents until the transaction
/// has had the opportunity to supply explicit replacement allocation facts.
pub(crate) fn delete_cascade_deferred(p: &mut Project, id: Id, pending: &mut BTreeSet<Id>) {
    let complete = deletion_impact(p, id);
    let removed = deletion_impact_inner(p, id, true);
    // An explicitly removed relation cannot pass its conditional cascade on
    // to a new relation that later reuses the same identity in this batch.
    pending.retain(|id| !removed.contains(id));
    pending.extend(
        complete
            .difference(&removed)
            .filter(|id| {
                p.relations
                    .get(id)
                    .is_some_and(|r| r.kind() == RelationKind::ExchangeAllocation)
            })
            .copied(),
    );
    remove_objects(p, removed);
}
pub(crate) fn finish_deletions(p: &mut Project, pending: &BTreeSet<Id>) {
    let none = BTreeSet::new();
    for id in pending {
        if p.relations.get(id).is_some_and(|r| {
            r.kind() == RelationKind::ExchangeAllocation
                && crate::ports::missing_exchange_port_allocations(p, r, &none)
        }) {
            delete_cascade(p, *id);
        }
    }
}
fn remove_objects(p: &mut Project, removed: BTreeSet<Id>) {
    let removed_names: std::collections::BTreeMap<_, _> = removed
        .iter()
        .map(|id| (*id, p.named(*id).to_string()))
        .collect();
    p.elements.retain(|id, _| !removed.contains(id));
    p.relations.retain(|id, _| !removed.contains(id));
    p.diagrams.retain(|id, _| !removed.contains(id));
    p.traces.retain(|id, _| !removed.contains(id));
    p.artifact_links.retain(|id, _| !removed.contains(id));
    for d in p.diagrams.values_mut() {
        d.nodes.retain(|id, _| !removed.contains(id));
    }
    let clear = |v: &mut Option<Id>| {
        if v.is_some_and(|id| removed.contains(&id)) {
            *v = None;
        }
    };
    for e in p.elements.values_mut() {
        if let Some(steps) = e.steps_mut() {
            steps.retain(|id| !removed.contains(id));
        }
        if let Some(shared) = e.shared_mut() {
            match shared {
                SharedBody::Interface { items } => items.retain(|id| !removed.contains(id)),
                SharedBody::Event {
                    body: EventBody::Signal { items },
                } => items.retain(|id| !removed.contains(id)),
                SharedBody::Port { contract, legacy } => {
                    match contract {
                        PortContract::FunctionInput { incoming_items } => {
                            incoming_items.retain(|id| !removed.contains(id))
                        }
                        PortContract::FunctionOutput { outgoing_items } => {
                            outgoing_items.retain(|id| !removed.contains(id))
                        }
                        _ => {}
                    }
                    if let Some(old) = legacy {
                        clear(&mut old.interface);
                    }
                }
                SharedBody::Scenario { context, .. } => clear(context),
                SharedBody::Lifeline { represented } => clear(represented),
                SharedBody::Constraint { constrained, .. } => {
                    constrained.retain(|id| !removed.contains(id))
                }
                SharedBody::Occurrence { body, .. } => match body {
                    OccurrenceBody::Exchange {
                        exchange,
                        sender,
                        receiver,
                        reply_to,
                        values,
                        ..
                    } => {
                        clear(exchange);
                        clear(sender);
                        clear(receiver);
                        clear(reply_to);
                        values.retain(|value| !removed.contains(&value.field));
                    }
                    OccurrenceBody::Action { action, lifeline } => {
                        clear(action);
                        clear(lifeline);
                    }
                    OccurrenceBody::SequenceReference { sequence } => clear(sequence),
                    OccurrenceBody::Control { .. } => {}
                    OccurrenceBody::ScenarioUse { scenario, bindings } => {
                        clear(scenario);
                        bindings.retain(|binding| {
                            !removed.contains(&binding.referenced)
                                && !removed.contains(&binding.local)
                        });
                    }
                    OccurrenceBody::State { state, lifeline } => {
                        clear(state);
                        clear(lifeline);
                    }
                    OccurrenceBody::Fragment { branches, .. } => {
                        branches.retain(|id| !removed.contains(id))
                    }
                    OccurrenceBody::Branch { .. } => {}
                },
                _ => {}
            }
        }
    }
    for r in p.relations.values_mut() {
        match &mut r.body {
            RelationBody::StateTransition {
                triggers,
                effects,
                trigger,
                ..
            } => {
                let deleted: Vec<_> = triggers
                    .iter()
                    .filter(|id| removed.contains(id))
                    .filter_map(|id| removed_names.get(id).cloned())
                    .collect();
                triggers.retain(|id| !removed.contains(id));
                effects.retain(|id| !removed.contains(id));
                if !deleted.is_empty() && triggers.is_empty() && trigger.trim().is_empty() {
                    *trigger = format!(
                        "Removed trigger: {}. Choose a replacement trigger.",
                        deleted.join(", ")
                    );
                }
            }
            RelationBody::Interaction { items, .. }
            | RelationBody::BoundaryExchange { items, .. } => {
                items.retain(|id| !removed.contains(id))
            }
            RelationBody::FunctionalExchange { items, .. } => {
                items.retain(|id| !removed.contains(id));
            }
            _ => {}
        }
    }
    // Removing an item or a containing payload field can disconnect a still-live
    // record field from an exchange. Its contextual annotation is no longer valid.
    let allowed: std::collections::BTreeMap<_, _> = p
        .elements
        .values()
        .filter_map(|e| match e.shared() {
            Some(SharedBody::Occurrence {
                body:
                    OccurrenceBody::Exchange {
                        exchange: Some(id), ..
                    },
                ..
            }) => p
                .relations
                .get(id)
                .map(|r| (e.id, crate::validation::carried_fields(p, r))),
            _ => None,
        })
        .collect();
    for (id, fields) in allowed {
        if let Some(SharedBody::Occurrence {
            body: OccurrenceBody::Exchange { values, .. },
            ..
        }) = p.elements.get_mut(&id).and_then(Element::shared_mut)
        {
            values.retain(|value| fields.contains(&value.field));
        }
    }
}