use std::collections::{BTreeMap, BTreeSet};

use super::char_support::get_char_value_from_grammar_char_literal;
use super::diagnostic::{CompilationError, Diagnostic};
use super::frontend::SourceFile;
use super::loader::LoadedSources;
use super::model::{
    Alternative, Authored, Block, ChannelDeclaration, Element, ElementId, ElementKind, GrammarId,
    GrammarKind, GrammarPrequel, GrammarUnit, Label, LexerCommand, Mode, ModelIdAllocator,
    ModelNodeId, NamedAction, Quantifier, Rule, RuleId, RuleKind, SetElement, Terminal,
    TokenDeclaration, VocabularySource,
};
use super::provenance::{MandatoryTransform, Origin, ProvenanceIndex, Tombstone};
use super::syntax::parse_grammar_unit;
use super::transform::TransformGrammar;
use super::validation::validate_model;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) struct RootOutputs {
    pub(crate) lexer: Option<GrammarId>,
    pub(crate) parser: Option<GrammarId>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) enum IntegratedVocabularySource {
    Grammar(GrammarId),
    TokensFile(std::path::PathBuf),
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) struct IntegratedVocabulary {
    pub(crate) consumer: GrammarId,
    pub(crate) source: IntegratedVocabularySource,
    pub(crate) declaration: Option<Authored<String>>,
}

#[derive(Clone, Debug)]
pub(crate) struct IntegratedGrammarSet {
    pub(crate) grammar: TransformGrammar,
    pub(crate) roots: BTreeMap<GrammarId, RootOutputs>,
    pub(crate) source_outputs: BTreeMap<GrammarId, RootOutputs>,
    pub(crate) vocabularies: Vec<IntegratedVocabulary>,
    pub(crate) diagnostics: Vec<Diagnostic>,
    pub(crate) ids: ModelIdAllocator,
}

pub(crate) fn integrate_loaded(
    loaded: &LoadedSources,
) -> Result<IntegratedGrammarSet, CompilationError> {
    let mut ids = ModelIdAllocator::after_loaded_grammars(loaded.grammars.grammars.len());
    let mut provenance = ProvenanceIndex::default();
    let source_units = loaded
        .grammars
        .grammars
        .iter()
        .enumerate()
        .map(|(index, parsed)| {
            let file = loaded
                .sources
                .get(parsed.source)
                .expect("loaded grammar source is retained");
            parse_grammar_unit(
                file,
                GrammarId::new(u32::try_from(index).expect("grammar count exceeds u32")),
                &mut ids,
                &mut provenance,
            )
        })
        .collect::<Vec<_>>();

    let mut diagnostics = loaded.diagnostics.clone();
    let targets = loaded
        .grammars
        .roots
        .iter()
        .copied()
        .chain(
            loaded
                .grammars
                .vocabularies
                .iter()
                .filter_map(|edge| match edge.source {
                    VocabularySource::Grammar(grammar) => Some(grammar),
                    VocabularySource::TokensFile(_) => None,
                }),
        )
        .collect::<BTreeSet<_>>();

    let mut units = Vec::new();
    let mut roots = BTreeMap::new();
    let mut source_outputs = BTreeMap::new();
    for source in loaded
        .grammars
        .load_order
        .iter()
        .copied()
        .filter(|grammar| targets.contains(grammar))
    {
        let mut unit = source_units[source.index()].clone();
        integrate_imports(
            source,
            &mut unit,
            &source_units,
            loaded,
            ImportIntegrationContext {
                ids: &mut ids,
                provenance: &mut provenance,
                diagnostics: &mut diagnostics,
            },
        );
        reduce_unit_blocks_to_sets(&mut unit, &mut ids, &mut provenance);
        if unit.kind == GrammarKind::Combined {
            let (lexer, parser) = split_combined(unit, &mut ids, &mut provenance);
            let lexer_id = lexer.as_ref().map(|unit| unit.id);
            let parser_id = parser.id;
            if let Some(lexer) = lexer {
                units.push(lexer);
            }
            units.push(parser);
            source_outputs.insert(
                source,
                RootOutputs {
                    lexer: lexer_id,
                    parser: Some(parser_id),
                },
            );
        } else {
            let outputs = match unit.kind {
                GrammarKind::Lexer => RootOutputs {
                    lexer: Some(unit.id),
                    parser: None,
                },
                GrammarKind::Parser => RootOutputs {
                    lexer: None,
                    parser: Some(unit.id),
                },
                GrammarKind::Combined => unreachable!("combined grammar handled above"),
            };
            source_outputs.insert(source, outputs);
            units.push(unit);
        }
    }

    for root in &loaded.grammars.roots {
        if let Some(outputs) = source_outputs.get(root).copied() {
            roots.insert(*root, outputs);
        }
    }

    let mut vocabularies = Vec::new();
    for (source, outputs) in &source_outputs {
        if loaded.grammars.grammar(*source).header.kind == GrammarKind::Combined {
            if let (Some(lexer), Some(parser)) = (outputs.lexer, outputs.parser) {
                vocabularies.push(IntegratedVocabulary {
                    consumer: parser,
                    source: IntegratedVocabularySource::Grammar(lexer),
                    declaration: None,
                });
            }
        }
    }
    for edge in &loaded.grammars.vocabularies {
        let Some(consumer_outputs) = source_outputs.get(&edge.importer) else {
            continue;
        };
        let consumer = consumer_outputs
            .parser
            .or(consumer_outputs.lexer)
            .expect("integrated source has an output grammar");
        let source = match &edge.source {
            VocabularySource::Grammar(source) => {
                let producer = source_outputs
                    .get(source)
                    .and_then(|outputs| outputs.lexer)
                    .expect("loader only accepts lexer-capable source vocabulary producers");
                IntegratedVocabularySource::Grammar(producer)
            }
            VocabularySource::TokensFile(path) => {
                IntegratedVocabularySource::TokensFile(path.clone())
            }
        };
        vocabularies.push(IntegratedVocabulary {
            consumer,
            source,
            declaration: Some(edge.declaration.clone()),
        });
    }

    let target_units = roots
        .values()
        .flat_map(|outputs| [outputs.lexer, outputs.parser])
        .flatten()
        .collect();
    let grammar = TransformGrammar {
        units,
        target_units,
        preserved_rules: BTreeSet::new(),
        provenance,
    };
    if let Err(diagnostic) = validate_model(&grammar) {
        return Err(CompilationError::new(vec![diagnostic]));
    }
    Ok(IntegratedGrammarSet {
        grammar,
        roots,
        source_outputs,
        vocabularies,
        diagnostics,
        ids,
    })
}

struct ImportIntegrationContext<'a> {
    ids: &'a mut ModelIdAllocator,
    provenance: &'a mut ProvenanceIndex,
    diagnostics: &'a mut Vec<Diagnostic>,
}

fn integrate_imports(
    root: GrammarId,
    destination: &mut GrammarUnit,
    source_units: &[GrammarUnit],
    loaded: &LoadedSources,
    context: ImportIntegrationContext<'_>,
) {
    let ImportIntegrationContext {
        ids,
        provenance,
        diagnostics,
    } = context;
    let (order, inbound_edges) = import_preorder(root, loaded);
    let mut rule_names = destination
        .rules
        .iter()
        .map(|rule| rule.name.clone())
        .collect::<BTreeSet<_>>();
    let mut channel_names = destination
        .channels
        .iter()
        .map(|channel| channel.name.value.clone())
        .collect::<BTreeSet<_>>();
    let mut action_keys = destination
        .actions
        .iter()
        .enumerate()
        .map(|(index, action)| (action_key(destination.kind, action), index))
        .collect::<BTreeMap<_, _>>();
    let mut mode_indices = destination
        .modes
        .iter()
        .enumerate()
        .map(|(index, mode)| (mode.name.clone(), index))
        .collect::<BTreeMap<_, _>>();
    let mut default_rule_insert = destination
        .rules
        .iter()
        .position(|rule| rule.mode.is_some())
        .unwrap_or(destination.rules.len());

    for imported_id in order {
        let imported = &source_units[imported_id.index()];
        let edges = inbound_edges
            .get(&imported_id)
            .map_or(&[][..], Vec::as_slice);
        let mut cloner = ImportedCloner {
            ids,
            provenance,
            edges,
        };

        for channel in &imported.channels {
            if channel_names.insert(channel.name.value.clone()) {
                destination.channels.push(cloner.channel(channel));
            }
        }
        destination
            .tokens
            .extend(imported.tokens.iter().map(|token| cloner.token(token)));

        for action in &imported.actions {
            let key = action_key(destination.kind, action);
            if let Some(index) = action_keys.get(&key).copied() {
                let existing = &mut destination.actions[index];
                if !existing.body.is_empty() && !action.body.is_empty() {
                    existing.body.push('\n');
                }
                existing.body.push_str(&action.body);
                cloner.merge_origin(
                    ModelNodeId::Action(existing.id),
                    ModelNodeId::Action(action.id),
                );
            } else {
                let cloned = cloner.action(action);
                action_keys.insert(key, destination.actions.len());
                destination.actions.push(cloned);
            }
        }

        for mode in &imported.modes {
            let surviving_rules = mode
                .rules
                .iter()
                .filter_map(|rule| {
                    imported
                        .rules
                        .iter()
                        .find(|candidate| candidate.id == *rule)
                })
                .filter(|rule| !rule_names.contains(&rule.name))
                .collect::<Vec<_>>();
            if surviving_rules.is_empty() {
                continue;
            }
            let destination_mode = if let Some(index) = mode_indices.get(&mode.name).copied() {
                destination.modes[index].id
            } else {
                let cloned = cloner.mode(mode);
                let id = cloned.id;
                mode_indices.insert(mode.name.clone(), destination.modes.len());
                destination.modes.push(cloned);
                id
            };
            for rule in surviving_rules {
                if rule_names.insert(rule.name.clone()) {
                    let cloned = cloner.rule(rule, Some(destination_mode));
                    let index = mode_indices[&mode.name];
                    destination.modes[index].rules.push(cloned.id);
                    destination.rules.push(cloned);
                }
            }
        }

        for rule in &imported.rules {
            if rule.mode.is_none() && rule_names.insert(rule.name.clone()) {
                destination
                    .rules
                    .insert(default_rule_insert, cloner.rule(rule, None));
                default_rule_insert += 1;
            }
        }

        if imported.options.iter().any(|option| {
            destination.options.iter().all(|root_option| {
                root_option.name.value != option.name.value
                    || root_option.value.value != option.value.value
            })
        }) {
            let primary = imported
                .prequels
                .iter()
                .find_map(|prequel| match prequel {
                    GrammarPrequel::Options { span, .. } => Some(span.clone()),
                    _ => None,
                })
                .unwrap_or_else(|| imported.span.clone());
            diagnostics.push(Diagnostic::warning(
                "G4T001",
                primary,
                format!(
                    "options from imported grammar {} do not alter root grammar {}",
                    imported.name, destination.name
                ),
            ));
        }
    }
}

fn import_preorder(
    root: GrammarId,
    loaded: &LoadedSources,
) -> (Vec<GrammarId>, BTreeMap<GrammarId, Vec<u32>>) {
    fn visit(
        grammar: GrammarId,
        loaded: &LoadedSources,
        seen: &mut BTreeSet<GrammarId>,
        order: &mut Vec<GrammarId>,
        inbound: &mut BTreeMap<GrammarId, Vec<u32>>,
    ) {
        for edge in loaded
            .grammars
            .imports
            .iter()
            .filter(|edge| edge.importer == grammar)
        {
            let edges = inbound.entry(edge.imported).or_default();
            if !edges.contains(&edge.id) {
                edges.push(edge.id);
            }
            if seen.insert(edge.imported) {
                order.push(edge.imported);
                visit(edge.imported, loaded, seen, order, inbound);
            }
        }
    }

    let mut seen = BTreeSet::from([root]);
    let mut order = Vec::new();
    let mut inbound = BTreeMap::new();
    visit(root, loaded, &mut seen, &mut order, &mut inbound);
    (order, inbound)
}

fn action_key(kind: GrammarKind, action: &NamedAction) -> (String, String) {
    let scope = action.scope.clone().unwrap_or_else(|| match kind {
        GrammarKind::Lexer => "lexer".to_owned(),
        GrammarKind::Parser | GrammarKind::Combined => "parser".to_owned(),
    });
    (scope, action.name.clone())
}

struct ImportedCloner<'a> {
    ids: &'a mut ModelIdAllocator,
    provenance: &'a mut ProvenanceIndex,
    edges: &'a [u32],
}

impl ImportedCloner<'_> {
    fn token(&mut self, source: &TokenDeclaration) -> TokenDeclaration {
        let mut cloned = source.clone();
        cloned.id = self.ids.token();
        self.record(ModelNodeId::Token(cloned.id), ModelNodeId::Token(source.id));
        cloned
    }

    fn channel(&mut self, source: &ChannelDeclaration) -> ChannelDeclaration {
        let mut cloned = source.clone();
        cloned.id = self.ids.channel();
        self.record(
            ModelNodeId::Channel(cloned.id),
            ModelNodeId::Channel(source.id),
        );
        cloned
    }

    fn action(&mut self, source: &NamedAction) -> NamedAction {
        let mut cloned = source.clone();
        cloned.id = self.ids.action();
        self.record(
            ModelNodeId::Action(cloned.id),
            ModelNodeId::Action(source.id),
        );
        cloned
    }

    fn mode(&mut self, source: &Mode) -> Mode {
        let mut cloned = source.clone();
        cloned.id = self.ids.mode();
        cloned.rules.clear();
        self.record(ModelNodeId::Mode(cloned.id), ModelNodeId::Mode(source.id));
        cloned
    }

    fn rule(&mut self, source: &Rule, mode: Option<super::model::ModeId>) -> Rule {
        let mut cloned = source.clone();
        cloned.id = self.ids.rule();
        cloned.mode = mode;
        cloned.actions = source
            .actions
            .iter()
            .map(|action| self.action(action))
            .collect();
        cloned.finally_action = source
            .finally_action
            .as_ref()
            .map(|action| self.action(action));
        cloned.block = self.block(&source.block);
        self.record(ModelNodeId::Rule(cloned.id), ModelNodeId::Rule(source.id));
        cloned
    }

    fn block(&mut self, source: &Block) -> Block {
        Block {
            alternatives: source
                .alternatives
                .iter()
                .map(|alternative| self.alternative(alternative))
                .collect(),
            options: source.options.clone(),
            syntax: source.syntax,
            span: source.span.clone(),
        }
    }

    fn alternative(&mut self, source: &Alternative) -> Alternative {
        let mut cloned = source.clone();
        cloned.id = self.ids.alternative();
        cloned.elements = source
            .elements
            .iter()
            .map(|element| self.element(element))
            .collect();
        self.record(
            ModelNodeId::Alternative(cloned.id),
            ModelNodeId::Alternative(source.id),
        );
        cloned
    }

    fn element(&mut self, source: &Element) -> Element {
        let mut cloned = source.clone();
        cloned.id = self.ids.element();
        cloned.label = source.label.as_ref().map(|label| self.label(label));
        cloned.kind = match &source.kind {
            ElementKind::Block(block) => ElementKind::Block(self.block(block)),
            ElementKind::Action { id, body } => {
                let cloned_id = self.ids.action();
                self.record(ModelNodeId::Action(cloned_id), ModelNodeId::Action(*id));
                ElementKind::Action {
                    id: cloned_id,
                    body: body.clone(),
                }
            }
            ElementKind::Predicate {
                id,
                body,
                fail,
                precedence,
            } => {
                let cloned_id = self.ids.predicate();
                self.record(
                    ModelNodeId::Predicate(cloned_id),
                    ModelNodeId::Predicate(*id),
                );
                ElementKind::Predicate {
                    id: cloned_id,
                    body: body.clone(),
                    fail: fail.clone(),
                    precedence: *precedence,
                }
            }
            kind => kind.clone(),
        };
        self.record(
            ModelNodeId::Element(cloned.id),
            ModelNodeId::Element(source.id),
        );
        cloned
    }

    fn label(&mut self, source: &Label) -> Label {
        let mut cloned = source.clone();
        cloned.id = self.ids.label();
        self.record(ModelNodeId::Label(cloned.id), ModelNodeId::Label(source.id));
        cloned
    }

    fn merge_origin(&mut self, destination: ModelNodeId, source: ModelNodeId) {
        let mut origins = self.provenance.origins(source).to_vec();
        origins.extend(self.edges.iter().map(|edge| Origin::Imported {
            edge: *edge,
            original: source,
        }));
        self.provenance.record_model(destination, origins);
    }

    fn record(&mut self, destination: ModelNodeId, source: ModelNodeId) {
        self.merge_origin(destination, source);
    }
}

fn reduce_unit_blocks_to_sets(
    unit: &mut GrammarUnit,
    ids: &mut ModelIdAllocator,
    provenance: &mut ProvenanceIndex,
) {
    for rule in &mut unit.rules {
        reduce_block_children(&mut rule.block, rule.kind == RuleKind::Lexer, provenance);
        reduce_top_level_block(
            &mut rule.block,
            rule.kind == RuleKind::Lexer,
            ids,
            provenance,
        );
    }
}

fn reduce_block_children(block: &mut Block, lexer: bool, provenance: &mut ProvenanceIndex) {
    for alternative in &mut block.alternatives {
        for element in &mut alternative.elements {
            let ElementKind::Block(nested) = &mut element.kind else {
                continue;
            };
            reduce_block_children(nested, lexer, provenance);
            if let Some((members, inputs, commands)) = block_set_members(nested, lexer) {
                if !commands.is_empty() {
                    continue;
                }
                let replacement = ModelNodeId::Element(element.id);
                record_transform(
                    provenance,
                    replacement,
                    MandatoryTransform::BlockSetReduction,
                    &inputs,
                );
                tombstone_block(provenance, nested, &[replacement]);
                element.kind = ElementKind::Set {
                    inverted: false,
                    elements: members,
                };
            }
        }
    }
}

fn reduce_top_level_block(
    block: &mut Block,
    lexer: bool,
    ids: &mut ModelIdAllocator,
    provenance: &mut ProvenanceIndex,
) {
    let Some((members, inputs, commands)) = block_set_members(block, lexer) else {
        return;
    };
    let new_alternative = ids.alternative();
    let new_element = ids.element();
    let old_alternatives = std::mem::take(&mut block.alternatives);
    let alternatives = old_alternatives
        .iter()
        .map(|alternative| ModelNodeId::Alternative(alternative.id))
        .collect::<Vec<_>>();
    record_transform(
        provenance,
        ModelNodeId::Alternative(new_alternative),
        MandatoryTransform::BlockSetReduction,
        &alternatives,
    );
    record_transform(
        provenance,
        ModelNodeId::Element(new_element),
        MandatoryTransform::BlockSetReduction,
        &inputs,
    );
    for alternative in &old_alternatives {
        provenance.tombstone(
            alternative.syntax,
            Tombstone {
                phase: "block-to-set",
                reason: "alternative merged into a set",
                replacements: Box::new([
                    ModelNodeId::Alternative(new_alternative),
                    ModelNodeId::Element(new_element),
                ]),
            },
        );
        for element in &alternative.elements {
            provenance.tombstone(
                element.syntax,
                Tombstone {
                    phase: "block-to-set",
                    reason: "element merged into a set",
                    replacements: Box::new([ModelNodeId::Element(new_element)]),
                },
            );
        }
    }
    block.alternatives.push(Alternative {
        id: new_alternative,
        elements: vec![Element {
            id: new_element,
            kind: ElementKind::Set {
                inverted: false,
                elements: members,
            },
            quantifier: Quantifier::One,
            label: None,
            options: Vec::new(),
            syntax: block.syntax,
            span: block.span.clone(),
            enclosing_span: block.span.clone(),
        }],
        label: None,
        options: Vec::new(),
        commands,
        syntax: block.syntax,
        span: block.span.clone(),
    });
}

fn block_set_members(
    block: &Block,
    lexer: bool,
) -> Option<(Vec<SetElement>, Vec<ModelNodeId>, Vec<LexerCommand>)> {
    if !block.options.is_empty() || block.alternatives.len() < 2 {
        return None;
    }
    let commands = block.alternatives.first()?.commands.clone();
    let mut members = Vec::with_capacity(block.alternatives.len());
    let mut inputs = Vec::with_capacity(block.alternatives.len());
    for alternative in &block.alternatives {
        if alternative.label.is_some()
            || !alternative.options.is_empty()
            || !same_lexer_commands(&alternative.commands, &commands)
            || alternative.elements.len() != 1
        {
            return None;
        }
        let element = &alternative.elements[0];
        if element.quantifier != Quantifier::One || element.label.is_some() {
            return None;
        }
        let member = match &element.kind {
            ElementKind::Terminal(Terminal::Literal(literal))
                if !lexer || grammar_literal_code_point(literal).is_some() =>
            {
                SetElement::Terminal {
                    source: element.id,
                    value: Terminal::Literal(literal.clone()),
                    span: element.span.clone(),
                    options: element.options.clone(),
                }
            }
            ElementKind::Terminal(Terminal::Token(token)) if !lexer => SetElement::Terminal {
                source: element.id,
                value: Terminal::Token(token.clone()),
                span: element.span.clone(),
                options: element.options.clone(),
            },
            ElementKind::Range(start, stop, operator_span)
                if lexer
                    && grammar_literal_code_point(&start.value)
                        .zip(grammar_literal_code_point(&stop.value))
                        .is_some_and(|(start, stop)| start <= stop) =>
            {
                SetElement::Range {
                    source: element.id,
                    start: start.value.clone(),
                    stop: stop.value.clone(),
                    span: operator_span.clone(),
                    options: element.options.clone(),
                }
            }
            _ => return None,
        };
        members.push(member);
        inputs.push(ModelNodeId::Element(element.id));
    }
    Some((members, inputs, commands))
}

fn same_lexer_commands(left: &[LexerCommand], right: &[LexerCommand]) -> bool {
    left.len() == right.len()
        && left
            .iter()
            .zip(right)
            .all(|(left, right)| left.name == right.name && left.argument == right.argument)
}

fn record_transform(
    provenance: &mut ProvenanceIndex,
    destination: ModelNodeId,
    transform: MandatoryTransform,
    inputs: &[ModelNodeId],
) {
    let mut origins = inputs
        .iter()
        .flat_map(|input| provenance.origins(*input).iter().cloned())
        .collect::<Vec<_>>();
    origins.push(Origin::MandatoryTransform {
        kind: transform,
        inputs: inputs.to_vec().into_boxed_slice(),
    });
    provenance.record_model(destination, origins);
}

fn tombstone_block(provenance: &mut ProvenanceIndex, block: &Block, replacements: &[ModelNodeId]) {
    for alternative in &block.alternatives {
        provenance.tombstone(
            alternative.syntax,
            Tombstone {
                phase: "block-to-set",
                reason: "nested alternative merged into a set",
                replacements: replacements.to_vec().into_boxed_slice(),
            },
        );
        for element in &alternative.elements {
            provenance.tombstone(
                element.syntax,
                Tombstone {
                    phase: "block-to-set",
                    reason: "nested element merged into a set",
                    replacements: replacements.to_vec().into_boxed_slice(),
                },
            );
        }
    }
}

fn grammar_literal_code_point(literal: &str) -> Option<i32> {
    let value = get_char_value_from_grammar_char_literal(Some(literal));
    (value >= 0).then_some(value)
}

fn split_combined(
    mut combined: GrammarUnit,
    ids: &mut ModelIdAllocator,
    provenance: &mut ProvenanceIndex,
) -> (Option<GrammarUnit>, GrammarUnit) {
    let combined_id = combined.id;
    let base_name = combined.name.clone();
    let modal_rules = combined
        .modes
        .iter()
        .flat_map(|mode| mode.rules.iter())
        .copied()
        .collect::<BTreeSet<_>>();
    let (parser_rules, mut lexer_rules) =
        partition_combined_rules(std::mem::take(&mut combined.rules), &modal_rules);
    combined.rules = parser_rules;

    let literals = implicit_literal_order(&combined.rules, &lexer_rules);
    let mut implicit_rules = Vec::new();
    for literal in literals {
        let (original, source) = find_literal_element(&combined.rules, &literal)
            .expect("collected parser literal has an owning element");
        implicit_rules.push(implicit_literal_rule(
            ImplicitLiteralSpec {
                name: format!("T__{}", implicit_rules.len()),
                literal,
                original,
                source,
                combined: combined_id,
            },
            ids,
            provenance,
        ));
    }
    implicit_rules.append(&mut lexer_rules);

    for rule in &implicit_rules {
        add_rule_implicit_origins(provenance, rule, combined_id);
    }

    let lexer_actions = combined
        .actions
        .iter()
        .map(|action| clone_implicit_action(action, combined_id, ids, provenance))
        .collect::<Vec<_>>();
    combined
        .actions
        .retain(|action| action.scope.as_deref() != Some("lexer"));

    let lexer_id = ids.grammar();
    let lexer = (!implicit_rules.is_empty()).then(|| {
        let mut origins = provenance
            .origins(ModelNodeId::Grammar(combined_id))
            .to_vec();
        origins.push(Origin::ImplicitLexer {
            combined: combined_id,
            original: ModelNodeId::Grammar(combined_id),
        });
        provenance.record_model(ModelNodeId::Grammar(lexer_id), origins);
        GrammarUnit {
            id: lexer_id,
            source: combined.source,
            name: format!("{base_name}Lexer"),
            kind: GrammarKind::Lexer,
            prequels: Vec::new(),
            options: combined
                .options
                .iter()
                .filter(|option| {
                    matches!(
                        option.name.value.as_str(),
                        "contextSuperClass"
                            | "language"
                            | "accessLevel"
                            | "exportMacro"
                            | "caseInsensitive"
                    )
                })
                .cloned()
                .collect(),
            tokens: combined
                .tokens
                .iter()
                .filter(|token| token.name.span.source == combined.source)
                .map(|token| clone_implicit_token(token, combined_id, ids, provenance))
                .collect(),
            channels: Vec::new(),
            actions: lexer_actions,
            modes: Vec::new(),
            rules: implicit_rules,
            syntax: combined.syntax,
            span: combined.span.clone(),
        }
    });

    combined.kind = GrammarKind::Parser;
    combined.name = format!("{base_name}Parser");
    (lexer, combined)
}

fn parser_literals(rules: &[Rule]) -> Vec<String> {
    fn visit(block: &Block, seen: &mut BTreeSet<String>, literals: &mut Vec<String>) {
        for alternative in &block.alternatives {
            for element in &alternative.elements {
                match &element.kind {
                    ElementKind::Terminal(Terminal::Literal(literal)) => {
                        if seen.insert(literal.clone()) {
                            literals.push(literal.clone());
                        }
                    }
                    ElementKind::Set { elements, .. } => {
                        for member in elements {
                            if let SetElement::Terminal {
                                value: Terminal::Literal(literal),
                                ..
                            } = member
                            {
                                if seen.insert(literal.clone()) {
                                    literals.push(literal.clone());
                                }
                            }
                        }
                    }
                    ElementKind::Block(nested) => visit(nested, seen, literals),
                    ElementKind::RuleCall(_)
                    | ElementKind::Terminal(_)
                    | ElementKind::Range(..)
                    | ElementKind::Action { .. }
                    | ElementKind::Predicate { .. }
                    | ElementKind::Epsilon => {}
                }
            }
        }
    }

    let mut seen = BTreeSet::new();
    let mut literals = Vec::new();
    for rule in rules {
        visit(&rule.block, &mut seen, &mut literals);
    }
    literals
}

pub(crate) fn source_implicit_token_literals(file: &SourceFile) -> Option<Vec<String>> {
    let mut ids = ModelIdAllocator::after_loaded_grammars(1);
    let mut provenance = ProvenanceIndex::default();
    let unit = parse_grammar_unit(file, GrammarId::new(0), &mut ids, &mut provenance);
    if unit.kind != GrammarKind::Combined {
        return None;
    }

    let modal_rules = unit
        .modes
        .iter()
        .flat_map(|mode| mode.rules.iter())
        .copied()
        .collect::<BTreeSet<_>>();
    let (parser_rules, lexer_rules) = partition_combined_rules(unit.rules, &modal_rules);
    Some(implicit_literal_order(&parser_rules, &lexer_rules))
}

fn partition_combined_rules(
    rules: Vec<Rule>,
    modal_rules: &BTreeSet<RuleId>,
) -> (Vec<Rule>, Vec<Rule>) {
    rules
        .into_iter()
        .partition(|rule| modal_rules.contains(&rule.id) || rule.kind == RuleKind::Parser)
}

fn implicit_literal_order(parser_rules: &[Rule], lexer_rules: &[Rule]) -> Vec<String> {
    let aliases = literal_aliases(lexer_rules);
    parser_literals(parser_rules)
        .into_iter()
        .filter(|literal| !aliases.contains(literal))
        .collect()
}

fn literal_aliases(rules: &[Rule]) -> BTreeSet<String> {
    rules
        .iter()
        .filter(|rule| !rule.fragment)
        .filter_map(|rule| {
            let alternative = rule.block.alternatives.first()?;
            if rule.block.alternatives.len() != 1 {
                return None;
            }
            let first = alternative.elements.first()?;
            if alternative.elements[1..].iter().any(|element| {
                !matches!(
                    element.kind,
                    ElementKind::Action { .. } | ElementKind::Predicate { .. }
                )
            }) {
                return None;
            }
            match &first.kind {
                ElementKind::Terminal(Terminal::Literal(literal))
                    if first.quantifier == Quantifier::One =>
                {
                    Some(literal.clone())
                }
                _ => None,
            }
        })
        .collect()
}

fn find_literal_element<'a>(rules: &'a [Rule], literal: &str) -> Option<(&'a Element, ElementId)> {
    fn find<'a>(block: &'a Block, literal: &str) -> Option<(&'a Element, ElementId)> {
        for alternative in &block.alternatives {
            for element in &alternative.elements {
                match &element.kind {
                    ElementKind::Terminal(Terminal::Literal(candidate)) if candidate == literal => {
                        return Some((element, element.id));
                    }
                    ElementKind::Set { elements, .. } => {
                        if let Some(source) = elements.iter().find_map(|member| match member {
                            SetElement::Terminal {
                                source,
                                value: Terminal::Literal(candidate),
                                ..
                            } if candidate == literal => Some(*source),
                            _ => None,
                        }) {
                            return Some((element, source));
                        }
                    }
                    ElementKind::Block(nested) => {
                        if let Some(found) = find(nested, literal) {
                            return Some(found);
                        }
                    }
                    _ => {}
                }
            }
        }
        None
    }
    rules.iter().find_map(|rule| find(&rule.block, literal))
}

struct ImplicitLiteralSpec<'a> {
    name: String,
    literal: String,
    original: &'a Element,
    source: ElementId,
    combined: GrammarId,
}

fn implicit_literal_rule(
    spec: ImplicitLiteralSpec<'_>,
    ids: &mut ModelIdAllocator,
    provenance: &mut ProvenanceIndex,
) -> Rule {
    let ImplicitLiteralSpec {
        name,
        literal,
        original,
        source,
        combined,
    } = spec;
    let rule_id = ids.rule();
    let alternative_id = ids.alternative();
    let element_id = ids.element();
    let original_id = ModelNodeId::Element(source);
    let mut origins = provenance.origins(original_id).to_vec();
    origins.push(Origin::ImplicitLexer {
        combined,
        original: original_id,
    });
    provenance.record_model(ModelNodeId::Rule(rule_id), origins.clone());
    provenance.record_model(ModelNodeId::Alternative(alternative_id), origins.clone());
    provenance.record_model(ModelNodeId::Element(element_id), origins);
    Rule {
        id: rule_id,
        name,
        name_span: original.span.clone(),
        kind: RuleKind::Lexer,
        fragment: false,
        modifiers: Vec::new(),
        arguments: None,
        returns: None,
        locals: None,
        throws: Vec::new(),
        options: Vec::new(),
        actions: Vec::new(),
        catches: Vec::new(),
        finally_action: None,
        left_recursion: None,
        block: Block {
            alternatives: vec![Alternative {
                id: alternative_id,
                elements: vec![Element {
                    id: element_id,
                    kind: ElementKind::Terminal(Terminal::Literal(literal)),
                    quantifier: Quantifier::One,
                    label: None,
                    options: Vec::new(),
                    syntax: original.syntax,
                    span: original.span.clone(),
                    enclosing_span: original.span.clone(),
                }],
                label: None,
                options: Vec::new(),
                commands: Vec::new(),
                syntax: original.syntax,
                span: original.span.clone(),
            }],
            options: Vec::new(),
            syntax: original.syntax,
            span: original.span.clone(),
        },
        mode: None,
        case_insensitive: None,
        syntax: original.syntax,
        span: original.span.clone(),
    }
}

fn clone_implicit_action(
    source: &NamedAction,
    combined: GrammarId,
    ids: &mut ModelIdAllocator,
    provenance: &mut ProvenanceIndex,
) -> NamedAction {
    let mut cloned = source.clone();
    cloned.id = ids.action();
    let source_id = ModelNodeId::Action(source.id);
    let mut origins = provenance.origins(source_id).to_vec();
    origins.push(Origin::ImplicitLexer {
        combined,
        original: source_id,
    });
    provenance.record_model(ModelNodeId::Action(cloned.id), origins);
    cloned
}

fn clone_implicit_token(
    source: &TokenDeclaration,
    combined: GrammarId,
    ids: &mut ModelIdAllocator,
    provenance: &mut ProvenanceIndex,
) -> TokenDeclaration {
    let mut cloned = source.clone();
    cloned.id = ids.token();
    let source_id = ModelNodeId::Token(source.id);
    let mut origins = provenance.origins(source_id).to_vec();
    origins.push(Origin::ImplicitLexer {
        combined,
        original: source_id,
    });
    provenance.record_model(ModelNodeId::Token(cloned.id), origins);
    cloned
}

fn add_rule_implicit_origins(provenance: &mut ProvenanceIndex, rule: &Rule, combined: GrammarId) {
    add_implicit_origin(
        provenance,
        ModelNodeId::Rule(rule.id),
        combined,
        ModelNodeId::Rule(rule.id),
    );
    for action in &rule.actions {
        add_implicit_origin(
            provenance,
            ModelNodeId::Action(action.id),
            combined,
            ModelNodeId::Action(action.id),
        );
    }
    if let Some(action) = &rule.finally_action {
        add_implicit_origin(
            provenance,
            ModelNodeId::Action(action.id),
            combined,
            ModelNodeId::Action(action.id),
        );
    }
    add_block_implicit_origins(provenance, &rule.block, combined);
}

fn add_block_implicit_origins(
    provenance: &mut ProvenanceIndex,
    block: &Block,
    combined: GrammarId,
) {
    for alternative in &block.alternatives {
        add_implicit_origin(
            provenance,
            ModelNodeId::Alternative(alternative.id),
            combined,
            ModelNodeId::Alternative(alternative.id),
        );
        for element in &alternative.elements {
            add_implicit_origin(
                provenance,
                ModelNodeId::Element(element.id),
                combined,
                ModelNodeId::Element(element.id),
            );
            if let Some(label) = &element.label {
                add_implicit_origin(
                    provenance,
                    ModelNodeId::Label(label.id),
                    combined,
                    ModelNodeId::Label(label.id),
                );
            }
            match &element.kind {
                ElementKind::Action { id, .. } => add_implicit_origin(
                    provenance,
                    ModelNodeId::Action(*id),
                    combined,
                    ModelNodeId::Action(*id),
                ),
                ElementKind::Predicate { id, .. } => add_implicit_origin(
                    provenance,
                    ModelNodeId::Predicate(*id),
                    combined,
                    ModelNodeId::Predicate(*id),
                ),
                ElementKind::Block(nested) => {
                    add_block_implicit_origins(provenance, nested, combined);
                }
                _ => {}
            }
        }
    }
}

fn add_implicit_origin(
    provenance: &mut ProvenanceIndex,
    destination: ModelNodeId,
    combined: GrammarId,
    original: ModelNodeId,
) {
    provenance.record_model(destination, [Origin::ImplicitLexer { combined, original }]);
}

#[cfg(test)]
mod tests {
    use std::fs;
    use std::io::Write as _;
    use std::path::PathBuf;

    use super::*;
    use crate::grammar::loader::{LoadOptions, load};
    use crate::grammar::semantics::analyze;
    use crate::grammar::transform::analysis::AnalysisInvalidation;
    use crate::grammar::transform::artifact::render_unmodified_sources;
    use crate::grammar::transform::{
        GrammarTransform, SafetyClass, TransformContext, TransformRegistry, TransformReport,
    };

    #[test]
    fn integrates_diamond_imports_reduces_sets_and_splits_combined_grammar() {
        let fixture = Fixture::new("integrate");
        fixture.write(
            "Root.g4",
            r#"
grammar Root;
import Left, Right;
@members { root(); }
root : 'if' | ID;
ID : 'id';
WS : ' ' -> skip;
"#,
        );
        fixture.write(
            "Left.g4",
            r#"
parser grammar Left;
import Shared;
@members { left(); }
root : 'override';
left : 'left';
"#,
        );
        fixture.write(
            "Right.g4",
            r#"
parser grammar Right;
import Shared;
@members { right(); }
right : 'right';
"#,
        );
        fixture.write("Shared.g4", "parser grammar Shared; shared : 'shared';");

        let loaded = load(LoadOptions {
            roots: vec![fixture.path("Root.g4")],
            library_directories: Vec::new(),
        })
        .expect("fixture should load");
        let integrated = integrate_loaded(&loaded).expect("fixture should integrate");
        let outputs = integrated.roots[&GrammarId::new(0)];
        let lexer = integrated
            .grammar
            .units
            .iter()
            .find(|unit| Some(unit.id) == outputs.lexer)
            .expect("implicit lexer");
        let parser = integrated
            .grammar
            .units
            .iter()
            .find(|unit| Some(unit.id) == outputs.parser)
            .expect("combined parser");

        assert_eq!(lexer.name, "RootLexer");
        assert_eq!(parser.name, "RootParser");
        assert_eq!(
            parser
                .rules
                .iter()
                .map(|rule| rule.name.as_str())
                .collect::<Vec<_>>(),
            ["root", "left", "shared", "right"]
        );
        assert!(matches!(
            parser.rules[0].block.alternatives[0].elements[0].kind,
            ElementKind::Set {
                inverted: false,
                ..
            }
        ));
        assert_eq!(parser.actions[0].body, " root(); \n left(); \n right(); ");
        assert_eq!(
            lexer
                .rules
                .iter()
                .map(|rule| rule.name.as_str())
                .collect::<Vec<_>>(),
            ["T__0", "T__1", "T__2", "T__3", "ID", "WS"]
        );

        let shared = parser
            .rules
            .iter()
            .find(|rule| rule.name == "shared")
            .expect("shared imported rule");
        let imported_edges = integrated
            .grammar
            .provenance
            .origins(ModelNodeId::Rule(shared.id))
            .iter()
            .filter_map(|origin| match origin {
                Origin::Imported { edge, .. } => Some(*edge),
                _ => None,
            })
            .collect::<BTreeSet<_>>();
        assert_eq!(imported_edges.len(), 2);
        validate_model(&integrated.grammar).expect("integrated model should be valid");
    }

    #[test]
    fn identical_lexer_commands_are_preserved_when_alternatives_reduce_to_a_set() {
        let fixture = Fixture::new("lexer-command-set");
        fixture.write(
            "Commands.g4",
            "lexer grammar Commands; A : 'a' -> skip | 'b' -> skip;",
        );

        let loaded = load(LoadOptions {
            roots: vec![fixture.path("Commands.g4")],
            library_directories: Vec::new(),
        })
        .expect("fixture should load");
        let integrated = integrate_loaded(&loaded).expect("fixture should integrate");
        let rule = &integrated.grammar.units[0].rules[0];

        let [alternative] = rule.block.alternatives.as_slice() else {
            panic!("identical command alternatives should reduce to one set");
        };
        assert!(matches!(
            alternative.elements[0].kind,
            ElementKind::Set { .. }
        ));
        assert_eq!(alternative.commands[0].name, "skip");
    }

    #[test]
    fn combined_split_copies_declared_tokens_to_the_implicit_lexer() {
        let fixture = Fixture::new("combined-tokens");
        fixture.write(
            "Root.g4",
            "grammar Root; tokens { INDENT } root : X; X : 'x' -> type(INDENT);",
        );

        let loaded = load(LoadOptions {
            roots: vec![fixture.path("Root.g4")],
            library_directories: Vec::new(),
        })
        .expect("fixture should load");
        let integrated = integrate_loaded(&loaded).expect("fixture should integrate");
        let lexer = integrated
            .grammar
            .units
            .iter()
            .find(|unit| unit.name == "RootLexer")
            .expect("implicit lexer");
        let parser = integrated
            .grammar
            .units
            .iter()
            .find(|unit| unit.name == "RootParser")
            .expect("combined parser");

        assert_eq!(lexer.tokens[0].name.value, "INDENT");
        assert_eq!(parser.tokens[0].name.value, "INDENT");
        assert_ne!(lexer.tokens[0].id, parser.tokens[0].id);
        let semantic = analyze(&loaded.sources, integrated)
            .expect("the implicit lexer type command should resolve");
        assert!(
            semantic
                .diagnostics
                .iter()
                .all(|diagnostic| diagnostic.code != "G4S019")
        );
    }

    #[test]
    fn fragment_literal_does_not_suppress_an_implicit_lexer_rule() {
        let fixture = Fixture::new("fragment-literal-alias");
        fixture.write("Root.g4", "grammar Root; root : 'x'; fragment F : 'x';");

        let loaded = load(LoadOptions {
            roots: vec![fixture.path("Root.g4")],
            library_directories: Vec::new(),
        })
        .expect("fixture should load");
        let integrated = integrate_loaded(&loaded).expect("fixture should integrate");
        let lexer = integrated
            .grammar
            .units
            .iter()
            .find(|unit| unit.name == "RootLexer")
            .expect("implicit lexer");

        assert_eq!(
            lexer
                .rules
                .iter()
                .map(|rule| rule.name.as_str())
                .collect::<Vec<_>>(),
            ["T__0", "F"]
        );
        analyze(&loaded.sources, integrated).expect("parser literal should resolve");
    }

    #[test]
    fn no_op_renderer_is_byte_identical() {
        let fixture = Fixture::new("render");
        let text = "parser grammar A;\n// retained trivia\na : 'x' ;\n";
        fixture.write("A.g4", text);
        let loaded = load(LoadOptions {
            roots: vec![fixture.path("A.g4")],
            library_directories: Vec::new(),
        })
        .expect("fixture should load");
        let rendered = render_unmodified_sources(&loaded.sources);
        let source = loaded
            .sources
            .iter()
            .next()
            .expect("fixture contains one source");
        assert_eq!(rendered[&source.id()], text);
    }

    #[test]
    fn integrates_dependency_only_source_vocabulary_before_its_consumer() {
        let fixture = Fixture::new("source-vocabulary");
        fixture.write(
            "Root.g4",
            "parser grammar Root; options { tokenVocab=Lex; } root : ID;",
        );
        fixture.write("Lex.g4", "lexer grammar Lex; ID : [a-z]+;");

        let loaded = load(LoadOptions {
            roots: vec![fixture.path("Root.g4")],
            library_directories: Vec::new(),
        })
        .expect("source-backed vocabulary should load");
        let integrated =
            integrate_loaded(&loaded).expect("source-backed vocabulary should integrate");

        assert_eq!(
            integrated
                .grammar
                .units
                .iter()
                .map(|unit| unit.name.as_str())
                .collect::<Vec<_>>(),
            ["Lex", "Root"]
        );
        assert_eq!(
            integrated.source_outputs[&GrammarId::new(1)].lexer,
            Some(GrammarId::new(1))
        );
        assert_eq!(
            integrated.vocabularies,
            [IntegratedVocabulary {
                consumer: GrammarId::new(0),
                source: IntegratedVocabularySource::Grammar(GrammarId::new(1)),
                declaration: loaded.grammars.vocabularies[0].declaration.clone().into(),
            }]
        );
    }

    #[test]
    fn imported_lexer_rules_precede_root_mode_rules() {
        let fixture = Fixture::new("lexer-import-modes");
        fixture.write(
            "Root.g4",
            "lexer grammar Root; import Base; A: 'a'; mode M; B: 'b';",
        );
        fixture.write(
            "Base.g4",
            "lexer grammar Base; fragment X: 'x'; fragment Y: 'y';",
        );

        let loaded = load(LoadOptions {
            roots: vec![fixture.path("Root.g4")],
            library_directories: Vec::new(),
        })
        .expect("fixture should load");
        let integrated = integrate_loaded(&loaded).expect("fixture should integrate");
        assert_eq!(
            integrated.grammar.units[0]
                .rules
                .iter()
                .map(|rule| rule.name.as_str())
                .collect::<Vec<_>>(),
            ["A", "X", "Y", "B"]
        );
    }

    #[test]
    fn mutation_recomputes_declared_analysis_before_the_next_pass() {
        struct Rename;
        impl GrammarTransform for Rename {
            fn name(&self) -> &'static str {
                "rename"
            }

            fn safety_class(&self) -> SafetyClass {
                SafetyClass::TreeAndApiPreserving
            }

            fn invalidates(&self) -> AnalysisInvalidation {
                AnalysisInvalidation::NAMES.union(AnalysisInvalidation::CALLS)
            }

            fn apply(
                &self,
                _input: &TransformContext<'_>,
                grammar: &mut TransformGrammar,
                _ids: &mut ModelIdAllocator,
                _report: &mut TransformReport,
            ) -> Result<bool, Diagnostic> {
                grammar.units[0].rules[0].name = "renamed".to_owned();
                Ok(true)
            }
        }

        struct Observe;
        impl GrammarTransform for Observe {
            fn name(&self) -> &'static str {
                "observe"
            }

            fn safety_class(&self) -> SafetyClass {
                SafetyClass::TreeAndApiPreserving
            }

            fn invalidates(&self) -> AnalysisInvalidation {
                AnalysisInvalidation::default()
            }

            fn apply(
                &self,
                input: &TransformContext<'_>,
                _grammar: &mut TransformGrammar,
                _ids: &mut ModelIdAllocator,
                _report: &mut TransformReport,
            ) -> Result<bool, Diagnostic> {
                assert!(input.analysis.rules_by_name.contains_key("renamed"));
                Ok(false)
            }
        }

        let fixture = Fixture::new("invalidate");
        fixture.write("A.g4", "parser grammar A; a : 'x';");
        let loaded = load(LoadOptions {
            roots: vec![fixture.path("A.g4")],
            library_directories: Vec::new(),
        })
        .expect("fixture should load");
        let mut integrated = integrate_loaded(&loaded).expect("fixture should integrate");
        let mut registry = TransformRegistry::default();
        registry.push(Rename);
        registry.push(Observe);
        let report = registry
            .run(&mut integrated.grammar, &mut integrated.ids, false)
            .expect("passes should run");
        assert_eq!(report.entries.len(), 2);
        assert!(report.entries[0].changed);
    }

    struct Fixture {
        root: PathBuf,
    }

    impl Fixture {
        fn new(name: &str) -> Self {
            static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
            let serial = NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            let root = std::env::temp_dir().join(format!(
                "antlr-rust-phase-b-transform-{name}-{}-{serial}",
                std::process::id()
            ));
            fs::create_dir_all(&root).expect("fixture directory");
            Self { root }
        }

        fn path(&self, relative: &str) -> PathBuf {
            self.root.join(relative)
        }

        fn write(&self, relative: &str, text: &str) {
            let path = self.path(relative);
            if let Some(parent) = path.parent() {
                fs::create_dir_all(parent).expect("fixture parent");
            }
            let mut file = fs::File::create(&path).expect("fixture file");
            file.write_all(text.as_bytes()).expect("fixture contents");
        }
    }

    impl Drop for Fixture {
        fn drop(&mut self) {
            let _ = fs::remove_dir_all(&self.root);
        }
    }
}
