use std::collections::BTreeMap;

use antlr4_runtime::atn::AtnStateKind;

use super::super::model::{BuildStateId, BuildTransitionId, ModelNodeId, RuleId};
use super::super::provenance::{Origin, ProvenanceIndex, SyntheticReason};

#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) enum BuildTransitionKind {
    Epsilon,
    Atom(i32),
    Range {
        start: i32,
        stop: i32,
    },
    Set(Vec<(i32, i32)>),
    NotSet(Vec<(i32, i32)>),
    Wildcard,
    Rule {
        rule: RuleId,
        rule_index: usize,
        follow: BuildStateId,
        precedence: i32,
    },
    Predicate {
        rule_index: usize,
        predicate_index: usize,
        context_dependent: bool,
    },
    Action {
        rule_index: usize,
        action_index: Option<usize>,
        context_dependent: bool,
    },
    Precedence(i32),
}

impl BuildTransitionKind {
    pub(crate) const fn is_epsilon(&self) -> bool {
        matches!(
            self,
            Self::Epsilon
                | Self::Rule { .. }
                | Self::Predicate { .. }
                | Self::Action { .. }
                | Self::Precedence(_)
        )
    }
}

#[derive(Clone, Debug)]
pub(crate) struct BuildTransition {
    pub(crate) id: BuildTransitionId,
    pub(crate) source: BuildStateId,
    pub(crate) target: BuildStateId,
    pub(crate) kind: BuildTransitionKind,
}

pub(crate) struct BuildTransitionSpec {
    pub(crate) source: BuildStateId,
    pub(crate) target: BuildStateId,
    pub(crate) kind: BuildTransitionKind,
    pub(crate) prepend: bool,
}

#[derive(Clone, Debug)]
pub(crate) struct BuildState {
    pub(crate) id: BuildStateId,
    pub(crate) kind: AtnStateKind,
    pub(crate) rule: Option<RuleId>,
    pub(crate) rule_index: Option<usize>,
    pub(crate) end_state: Option<BuildStateId>,
    pub(crate) loop_back_state: Option<BuildStateId>,
    pub(crate) non_greedy: bool,
    pub(crate) left_recursive_rule: bool,
    pub(crate) transitions: Vec<BuildTransitionId>,
    pub(crate) active: bool,
}

#[derive(Clone, Debug)]
pub(crate) struct BuildGraph {
    pub(crate) max_token_type: i32,
    pub(crate) states: Vec<BuildState>,
    pub(crate) transitions: Vec<BuildTransition>,
    pub(crate) decisions: Vec<BuildStateId>,
    pub(crate) rule_starts: Vec<BuildStateId>,
    pub(crate) rule_stops: Vec<BuildStateId>,
    pub(crate) model_transitions: BTreeMap<ModelNodeId, Vec<BuildTransitionId>>,
}

impl BuildGraph {
    pub(crate) const fn new(max_token_type: i32) -> Self {
        Self {
            max_token_type,
            states: Vec::new(),
            transitions: Vec::new(),
            decisions: Vec::new(),
            rule_starts: Vec::new(),
            rule_stops: Vec::new(),
            model_transitions: BTreeMap::new(),
        }
    }

    pub(crate) fn add_state(
        &mut self,
        kind: AtnStateKind,
        rule: Option<(RuleId, usize)>,
        origins: impl IntoIterator<Item = Origin>,
        provenance: &mut ProvenanceIndex,
    ) -> BuildStateId {
        let id = BuildStateId::new(
            u32::try_from(self.states.len()).expect("ATN state count exceeds u32"),
        );
        let (rule, rule_index) =
            rule.map_or((None, None), |(rule, index)| (Some(rule), Some(index)));
        self.states.push(BuildState {
            id,
            kind,
            rule,
            rule_index,
            end_state: None,
            loop_back_state: None,
            non_greedy: false,
            left_recursive_rule: false,
            transitions: Vec::new(),
            active: true,
        });
        provenance.record_state(id, origins);
        id
    }

    pub(crate) fn add_synthetic_state(
        &mut self,
        kind: AtnStateKind,
        rule: Option<(RuleId, usize)>,
        reason: SyntheticReason,
        owner: ModelNodeId,
        provenance: &mut ProvenanceIndex,
    ) -> BuildStateId {
        let mut origins = provenance.origins(owner).to_vec();
        origins.push(Origin::Synthetic { reason, owner });
        self.add_state(kind, rule, origins, provenance)
    }

    pub(crate) fn add_transition(
        &mut self,
        spec: BuildTransitionSpec,
        origins: impl IntoIterator<Item = Origin>,
        provenance: &mut ProvenanceIndex,
    ) -> BuildTransitionId {
        let BuildTransitionSpec {
            source,
            target,
            kind,
            prepend,
        } = spec;
        if let Some(existing) = self.state(source).transitions.iter().copied().find(|id| {
            let transition = self.transition(*id);
            transition.target == target && transition.kind == kind
        }) {
            provenance.record_transition(existing, origins);
            return existing;
        }
        let id = BuildTransitionId::new(
            u32::try_from(self.transitions.len()).expect("ATN transition count exceeds u32"),
        );
        self.transitions.push(BuildTransition {
            id,
            source,
            target,
            kind,
        });
        let transitions = &mut self.state_mut(source).transitions;
        if prepend {
            transitions.insert(0, id);
        } else {
            transitions.push(id);
        }
        provenance.record_transition(id, origins);
        id
    }

    pub(crate) fn add_synthetic_transition(
        &mut self,
        spec: BuildTransitionSpec,
        reason: SyntheticReason,
        owner: ModelNodeId,
        provenance: &mut ProvenanceIndex,
    ) -> BuildTransitionId {
        let mut origins = provenance.origins(owner).to_vec();
        origins.push(Origin::Synthetic { reason, owner });
        let transition = self.add_transition(spec, origins, provenance);
        self.record_model_transition(owner, transition);
        transition
    }

    pub(crate) fn record_model_transition(
        &mut self,
        owner: ModelNodeId,
        transition: BuildTransitionId,
    ) {
        let transitions = self.model_transitions.entry(owner).or_default();
        if !transitions.contains(&transition) {
            transitions.push(transition);
        }
    }

    pub(crate) fn state(&self, id: BuildStateId) -> &BuildState {
        &self.states[id.index()]
    }

    pub(crate) fn state_mut(&mut self, id: BuildStateId) -> &mut BuildState {
        &mut self.states[id.index()]
    }

    pub(crate) fn transition(&self, id: BuildTransitionId) -> &BuildTransition {
        &self.transitions[id.index()]
    }

    pub(crate) fn transition_mut(&mut self, id: BuildTransitionId) -> &mut BuildTransition {
        &mut self.transitions[id.index()]
    }

    pub(crate) fn remove_state(&mut self, id: BuildStateId) {
        self.state_mut(id).active = false;
    }

    pub(crate) fn add_decision(&mut self, state: BuildStateId) {
        self.decisions.push(state);
    }

    pub(crate) fn finalize(self) -> FinalizedAtnGraph {
        let mut state_map = BTreeMap::new();
        for state in self.states.iter().filter(|state| state.active) {
            let compact = state_map.len();
            state_map.insert(state.id, compact);
        }
        let transition_map = self
            .transitions
            .iter()
            .filter(|transition| {
                state_map.contains_key(&transition.source)
                    && state_map.contains_key(&transition.target)
            })
            .enumerate()
            .map(|(index, transition)| (transition.id, index))
            .collect::<BTreeMap<_, _>>();
        let model_transitions = self
            .model_transitions
            .into_iter()
            .filter_map(|(model, transitions)| {
                let transitions = transitions
                    .into_iter()
                    .filter_map(|transition| transition_map.get(&transition).copied())
                    .collect::<Vec<_>>();
                (!transitions.is_empty()).then_some((model, transitions))
            })
            .collect();

        let states = self
            .states
            .into_iter()
            .filter(|state| state.active)
            .map(|state| {
                let transitions = state
                    .transitions
                    .into_iter()
                    .filter(|transition| {
                        let transition = &self.transitions[transition.index()];
                        state_map.contains_key(&transition.target)
                    })
                    .collect();
                FinalizedState {
                    original: state.id,
                    kind: state.kind,
                    rule: state.rule,
                    rule_index: state.rule_index,
                    end_state: state
                        .end_state
                        .and_then(|target| state_map.get(&target).copied()),
                    loop_back_state: state
                        .loop_back_state
                        .and_then(|target| state_map.get(&target).copied()),
                    non_greedy: state.non_greedy,
                    left_recursive_rule: state.left_recursive_rule,
                    transitions,
                }
            })
            .collect();
        let transitions = self
            .transitions
            .into_iter()
            .filter(|transition| {
                state_map.contains_key(&transition.source)
                    && state_map.contains_key(&transition.target)
            })
            .map(|transition| FinalizedTransition {
                original: transition.id,
                source: state_map[&transition.source],
                target: state_map[&transition.target],
                kind: finalize_transition_kind(transition.kind, &state_map),
            })
            .collect();
        FinalizedAtnGraph {
            max_token_type: self.max_token_type,
            states,
            transitions,
            decisions: self
                .decisions
                .into_iter()
                .filter_map(|state| state_map.get(&state).copied())
                .collect(),
            rule_starts: self
                .rule_starts
                .into_iter()
                .map(|state| state_map[&state])
                .collect(),
            rule_stops: self
                .rule_stops
                .into_iter()
                .map(|state| state_map[&state])
                .collect(),
            state_map,
            model_transitions,
        }
    }
}

fn finalize_transition_kind(
    kind: BuildTransitionKind,
    states: &BTreeMap<BuildStateId, usize>,
) -> FinalizedTransitionKind {
    match kind {
        BuildTransitionKind::Epsilon => FinalizedTransitionKind::Epsilon,
        BuildTransitionKind::Atom(label) => FinalizedTransitionKind::Atom(label),
        BuildTransitionKind::Range { start, stop } => {
            FinalizedTransitionKind::Range { start, stop }
        }
        BuildTransitionKind::Set(ranges) => FinalizedTransitionKind::Set(ranges),
        BuildTransitionKind::NotSet(ranges) => FinalizedTransitionKind::NotSet(ranges),
        BuildTransitionKind::Wildcard => FinalizedTransitionKind::Wildcard,
        BuildTransitionKind::Rule {
            rule,
            rule_index,
            follow,
            precedence,
        } => FinalizedTransitionKind::Rule {
            rule,
            rule_index,
            follow: states[&follow],
            precedence,
        },
        BuildTransitionKind::Predicate {
            rule_index,
            predicate_index,
            context_dependent,
        } => FinalizedTransitionKind::Predicate {
            rule_index,
            predicate_index,
            context_dependent,
        },
        BuildTransitionKind::Action {
            rule_index,
            action_index,
            context_dependent,
        } => FinalizedTransitionKind::Action {
            rule_index,
            action_index,
            context_dependent,
        },
        BuildTransitionKind::Precedence(precedence) => {
            FinalizedTransitionKind::Precedence(precedence)
        }
    }
}

#[derive(Clone, Debug)]
pub(crate) struct FinalizedAtnGraph {
    pub(crate) max_token_type: i32,
    pub(crate) states: Vec<FinalizedState>,
    pub(crate) transitions: Vec<FinalizedTransition>,
    pub(crate) decisions: Vec<usize>,
    pub(crate) rule_starts: Vec<usize>,
    pub(crate) rule_stops: Vec<usize>,
    pub(crate) state_map: BTreeMap<BuildStateId, usize>,
    pub(crate) model_transitions: BTreeMap<ModelNodeId, Vec<usize>>,
}

impl FinalizedAtnGraph {
    pub(crate) fn transitions_for_model(
        &self,
        model: ModelNodeId,
    ) -> impl Iterator<Item = &FinalizedTransition> {
        self.model_transitions
            .get(&model)
            .into_iter()
            .flatten()
            .filter_map(|index| self.transitions.get(*index))
    }
}

#[derive(Clone, Debug)]
pub(crate) struct FinalizedState {
    pub(crate) original: BuildStateId,
    pub(crate) kind: AtnStateKind,
    pub(crate) rule: Option<RuleId>,
    pub(crate) rule_index: Option<usize>,
    pub(crate) end_state: Option<usize>,
    pub(crate) loop_back_state: Option<usize>,
    pub(crate) non_greedy: bool,
    pub(crate) left_recursive_rule: bool,
    pub(crate) transitions: Vec<BuildTransitionId>,
}

#[derive(Clone, Debug)]
pub(crate) struct FinalizedTransition {
    pub(crate) original: BuildTransitionId,
    pub(crate) source: usize,
    pub(crate) target: usize,
    pub(crate) kind: FinalizedTransitionKind,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) enum FinalizedTransitionKind {
    Epsilon,
    Atom(i32),
    Range {
        start: i32,
        stop: i32,
    },
    Set(Vec<(i32, i32)>),
    NotSet(Vec<(i32, i32)>),
    Wildcard,
    Rule {
        rule: RuleId,
        rule_index: usize,
        follow: usize,
        precedence: i32,
    },
    Predicate {
        rule_index: usize,
        predicate_index: usize,
        context_dependent: bool,
    },
    Action {
        rule_index: usize,
        action_index: Option<usize>,
        context_dependent: bool,
    },
    Precedence(i32),
}
