use std::{
    collections::HashSet,
    hash::{DefaultHasher, Hash, Hasher},
};

pub trait Stage: Clone + std::fmt::Debug + std::hash::Hash + Eq {
    type Consume: Clone + std::fmt::Debug + std::hash::Hash + Eq;
}

#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub struct Epsilon;
impl Stage for Epsilon {
    type Consume = Option<ByteRange>;
}

#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub struct Resolved;
impl Stage for Resolved {
    type Consume = ByteRange;
}

use crate::parse::regex::{LookDirection, LookPolarity};

use super::Pattern;
use super::byte_range::ByteRange;

/// NFA with epsilon transitions
#[derive(Clone)]
#[allow(clippy::upper_case_acronyms)]
pub struct ENFA<S: Stage> {
    pub states: Vec<EState<S>>,
}

#[derive(Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
struct Thread {
    own_state: StateId,
    positives: Vec<Thread>,
    negatives: Vec<Thread>,
}

fn cartesian_product<T: Clone>(x: Vec<Vec<T>>) -> Vec<Vec<T>> {
    let mut result = vec![Vec::new()];

    for xs in x {
        let mut next = Vec::new();

        for prefix in &result {
            for item in &xs {
                let mut v = prefix.clone();
                v.push(item.clone());
                next.push(v);
            }
        }

        result = next;
    }

    result
}

#[cfg(test)]
mod product_tests {
    use super::cartesian_product;

    #[test]
    fn basic_case() {
        let x = vec![vec![1, 2], vec![10, 20]];

        let out = cartesian_product(x);

        assert_eq!(
            out,
            vec![vec![1, 10], vec![1, 20], vec![2, 10], vec![2, 20],]
        );
    }

    #[test]
    fn single_dimension() {
        let x = vec![vec![1, 2, 3]];

        let out = cartesian_product(x);

        assert_eq!(out, vec![vec![1], vec![2], vec![3],]);
    }

    #[test]
    fn empty_outer_vector() {
        let x: Vec<Vec<i32>> = vec![];

        let out = cartesian_product(x);

        // One empty combination.
        assert_eq!(out, vec![vec![]]);
    }

    #[test]
    fn empty_inner_vector() {
        let x = vec![vec![1, 2], vec![], vec![3, 4]];

        let out = cartesian_product(x);

        assert!(out.is_empty());
    }

    #[test]
    fn output_size_matches_product() {
        let x = vec![vec![1, 2, 3], vec![4, 5], vec![6, 7, 8, 9]];

        let out = cartesian_product(x);

        assert_eq!(out.len(), 3 * 2 * 4);
    }

    #[test]
    fn every_output_has_correct_length() {
        let x = vec![vec!['a', 'b'], vec!['x', 'y', 'z'], vec!['0']];

        let out = cartesian_product(x);

        assert!(out.iter().all(|v| v.len() == 3));
    }

    #[test]
    fn works_with_strings() {
        let x = vec![
            vec!["a".to_string(), "b".to_string()],
            vec!["x".to_string()],
        ];

        let out = cartesian_product(x);

        assert_eq!(
            out,
            vec![
                vec!["a".to_string(), "x".to_string()],
                vec!["b".to_string(), "x".to_string()],
            ]
        );
    }
}

impl Thread {
    fn new(
        enfa: &ENFA<Resolved>,
        own_state: StateId,
        mut positives: Vec<Thread>,
        mut negatives: Vec<Thread>,
    ) -> Option<Self> {
        positives.sort();
        positives.dedup();
        positives.retain(|t| !t.accept(enfa));
        positives.shrink_to_fit();

        negatives.sort();
        negatives.dedup();
        negatives.shrink_to_fit();

        if negatives.iter().any(|t| t.accept(enfa)) {
            return None;
        }

        Some(Self {
            own_state,
            positives,
            negatives,
        })
    }

    fn accept(&self, enfa: &ENFA<Resolved>) -> bool {
        let pos = self.positives.iter().all(|t| t.accept(enfa));
        let neg = self.negatives.iter().all(|t| !t.accept(enfa));
        let this = match enfa.states[self.own_state].accept {
            Acceptance::Accept => true,
            Acceptance::Assertion => true,
            Acceptance::NotYet => false,
        };
        pos && neg && this
    }

    fn possible_transitions(&self, enfa: &ENFA<Resolved>, f: &mut impl FnMut(ByteRange)) {
        for t in enfa.states[self.own_state].trans.iter() {
            f(t.consumes);
        }
        for t in self.positives.iter() {
            t.possible_transitions(enfa, f);
        }
        for t in self.negatives.iter() {
            t.possible_transitions(enfa, f);
        }
    }

    fn transition(&self, enfa: &ENFA<Resolved>, ch: ByteRange) -> Vec<Self> {
        // if any of our negatives currently accept, we must kill this thread rn.
        if self.negatives.iter().any(|t| t.accept(enfa)) {
            return Vec::new();
        }

        let pos: Vec<Vec<Self>> = self
            .positives
            .iter()
            .map(|t| t.transition(enfa, ch))
            .collect();
        let neg: Vec<Vec<Self>> = self
            .negatives
            .iter()
            .map(|t| t.transition(enfa, ch))
            .collect();

        let pos = cartesian_product(pos);
        let neg = cartesian_product(neg);

        let mut out = Vec::new();

        for p in pos {
            for n in neg.clone() {
                for transition in enfa.states[self.own_state].trans.iter() {
                    if transition.consumes.overlaps(ch) {
                        let mut p = p.clone();
                        let mut n = n.clone();
                        for a in transition.assert.iter() {
                            if let Some(thread) = Thread::new(enfa, a.to, Vec::new(), Vec::new()) {
                                match a.polarity {
                                    LookPolarity::Positive => p.push(thread),
                                    LookPolarity::Negative => n.push(thread),
                                }
                            }
                        }
                        if let Some(thread) = Thread::new(enfa, transition.to, p, n) {
                            out.push(thread);
                        }
                    }
                }
            }
        }

        out
    }
}

#[derive(Clone)]
pub struct MultiState<'a> {
    nfa: &'a ENFA<Resolved>,
    threads: Vec<Thread>,
    accept: bool,
    hash: u64,
}

impl<'a> PartialEq for MultiState<'a> {
    fn eq(&self, other: &Self) -> bool {
        (self.nfa as *const ENFA<Resolved> as u64) == (other.nfa as *const ENFA<Resolved> as u64)
            && self.threads == other.threads
            && self.accept == other.accept
            && self.hash == other.hash
    }
}
impl<'a> Eq for MultiState<'a> {}

impl<'a> MultiState<'a> {
    fn new(nfa: &'a ENFA<Resolved>, mut threads: Vec<Thread>) -> Self {
        threads.sort();
        threads.dedup();
        threads.shrink_to_fit();

        let accept = threads.iter().any(|t| t.accept(nfa));
        let mut hasher = DefaultHasher::new();
        threads.hash(&mut hasher);
        let hash = hasher.finish();

        Self {
            nfa,
            threads,
            accept,
            hash,
        }
    }

    /// all the chars that will make an interesting transition
    pub fn possible_transitions(&self) -> Vec<ByteRange> {
        let mut vec = Vec::new();
        for t in self.threads.iter() {
            t.possible_transitions(self.nfa, &mut |x| vec.push(x));
        }
        vec = ByteRange::non_overlapping(vec);
        vec.sort();
        vec.dedup();
        vec.shrink_to_fit();
        vec
    }

    pub fn transition(&self, ch: ByteRange) -> Self {
        let new_states = self
            .threads
            .iter()
            .flat_map(|t| t.transition(self.nfa, ch))
            .collect();

        Self::new(self.nfa, new_states)
    }

    pub fn accept(&self) -> bool {
        self.accept
    }
}

impl<'a> Hash for MultiState<'a> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.hash.hash(state)
    }
}

macro_rules! set {
    () => {
        std::collections::HashSet::new()
    };
    ( $( $x:expr ),* ) => {{
        let mut set = std::collections::HashSet::new();
        $(
            set.insert($x);
        )*
        set
    }};
}

impl<S: Stage> ENFA<S> {
    fn shift(self, amt: usize) -> Vec<EState<S>> {
        let mut s = self.states;

        for state in s.iter_mut() {
            state.remap(|i| i + amt);
            if state.accept == Acceptance::Accept {
                state.accept = Acceptance::NotYet;
            }
        }

        s
    }

    pub fn remove_unreachable(&mut self) {
        let mut used = vec![false; self.states.len()];
        used[0] = true;
        for s in self.states.iter() {
            for i in s.reachable_states() {
                used[i] = true;
            }
        }
        let mut remap = vec![0; self.states.len()];
        let mut shift = 0;
        for i in 0..self.states.len() {
            if used[i] {
                remap[i] = i - shift;
            } else {
                shift += 1;
            }
        }
        for i in (0..self.states.len()).rev() {
            if !used[i] {
                self.states.remove(i);
            }
        }
        for s in self.states.iter_mut() {
            s.remap(|i| remap[i]);
        }
    }
}

impl ENFA<Epsilon> {
    fn epsilon_dfs(
        &self,
        i: StateId,
        visited: &mut [Option<Vec<Assertion>>],
        path: Vec<Assertion>,
    ) {
        if visited[i].is_some() {
            return;
        }
        visited[i] = Some(path.clone());
        for t in self.states[i].trans.iter() {
            if t.is_epsilon() {
                let mut path = path.clone();
                path.append(&mut t.assert.clone());
                self.epsilon_dfs(t.to, visited, path);
            }
        }
    }

    pub fn resolve_epsilon(self) -> ENFA<Resolved> {
        // state X --> { state Y, Z, W which get inlined into X }
        let includes: Vec<Vec<(StateId, Vec<Assertion>)>> = (0..self.states.len())
            .map(|i| {
                let mut reach = vec![None; self.states.len()];
                self.epsilon_dfs(i, &mut reach, Vec::new());
                reach
                    .into_iter()
                    .enumerate()
                    .filter_map(|(x, r)| r.map(|r| (x, r)))
                    .collect()
            })
            .collect();

        // states without epsilon transitions
        let mut states: Vec<EState<Resolved>> = self
            .states
            .into_iter()
            .map(EState::remove_epsilon)
            .collect();

        // inline real transitions
        for i in 0..states.len() {
            for (k, assertions) in includes[i].iter() {
                let new: Vec<Transition<_>> = states[*k]
                    .trans
                    .iter()
                    .cloned()
                    .map(|t| t.with_assertions(assertions.clone()))
                    .collect();
                states[i].trans.extend(new);
                if states[*k].accept != Acceptance::NotYet {
                    states[i].accept = states[*k].accept;
                }
            }
        }

        ENFA { states }
    }

    fn looping(self) -> Self {
        let mut states = vec![EState::start()];
        states.append(&mut self.shift(1));
        let len = states.len();
        states[0].set_epsilon_transitions([Transition::epsilon(1), Transition::epsilon(len)]);
        states[len - 1].set_epsilon_transitions([Transition::epsilon(0), Transition::epsilon(len)]);
        states.push(EState::terminal());
        Self { states }
    }

    fn repeat(self, times: usize) -> Self {
        let reps = vec![self; times];
        Self::concat(reps)
    }

    /// between 0 and x repetitions
    fn optx(self, x: usize) -> Self {
        let len = self.states.len();
        let mut repped = self.repeat(x);
        assert_eq!(repped.states.len(), x * len);
        for i in 1..=x {
            repped.states[0]
                .trans
                .insert(Transition::epsilon(i * len - 1));
        }
        repped
    }

    fn concat(nfas: Vec<Self>) -> Self {
        if nfas.is_empty() {
            return Self {
                states: vec![EState::terminal()],
            };
        }

        let mut states: Vec<EState<Epsilon>> = Vec::new();
        for nfa in nfas.into_iter() {
            let len = states.len();
            let mut ns = nfa.shift(len);
            if let Some(n) = states.last_mut() {
                n.trans.retain(|t| t.consumes.is_some());
                n.trans.insert(Transition::epsilon(len));
            }
            states.append(&mut ns);
        }

        let len = states.len();
        states[len - 1].accept = Acceptance::Accept;

        Self { states }
    }
}

impl ENFA<Resolved> {
    pub fn start_multi_state<'a>(&'a self) -> MultiState<'a> {
        MultiState::new(
            self,
            vec![Thread::new(self, 0, Vec::new(), Vec::new()).unwrap()],
        )
    }

    pub fn void_multi_state<'a>(&'a self) -> MultiState<'a> {
        MultiState::new(self, vec![])
    }

    pub fn all_multi_states<'a>(&'a self) -> HashSet<MultiState<'a>> {
        let mut states = set![self.start_multi_state()];
        let mut q = vec![self.start_multi_state()];

        while let Some(state) = q.pop() {
            let chars = state.possible_transitions();

            for chr in chars {
                let new = state.transition(chr);

                if !states.contains(&new) {
                    states.insert(new.clone());
                    q.push(new);
                }
            }
        }

        states
    }
}

impl std::fmt::Debug for ENFA<Epsilon> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "NFA {{")?;
        for (i, s) in self.states.iter().enumerate() {
            write!(f, "  {i}: ")?;
            for t in s.trans.iter() {
                let k = t.to;
                if let Some(c) = t.consumes {
                    write!(f, "{c:?}=>{k} ")?;
                } else {
                    write!(f, "~>{k} ")?;
                }
            }
            match s.accept {
                Acceptance::Accept => write!(f, "accept")?,
                Acceptance::Assertion => write!(f, "assert")?,
                Acceptance::NotYet => {}
            }
            writeln!(f)?;
        }
        write!(f, "}}")
    }
}

pub type StateId = usize;

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct Assertion {
    to: StateId,
    polarity: LookPolarity,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct Transition<S: Stage> {
    to: StateId,
    assert: Vec<Assertion>,
    consumes: S::Consume,
}

impl Transition<Epsilon> {
    fn new(consumes: ByteRange, to: StateId) -> Self {
        Self {
            to,
            consumes: Some(consumes),
            assert: Vec::new(),
        }
    }

    fn epsilon(to: StateId) -> Self {
        Self {
            to,
            consumes: None,
            assert: Vec::new(),
        }
    }

    pub fn is_epsilon(&self) -> bool {
        self.consumes.is_none()
    }

    fn non_epsilon(self) -> Option<Transition<Resolved>> {
        let Self {
            to,
            assert,
            consumes,
        } = self;
        let consumes = consumes?;
        Some(Transition {
            consumes,
            to,
            assert,
        })
    }
}

impl<S: Stage> Transition<S> {
    fn remap(&mut self, mut f: impl FnMut(StateId) -> StateId) {
        self.to = f(self.to);
    }

    fn reachable_states(&self) -> impl Iterator<Item = StateId> {
        [self.to].into_iter()
    }

    fn with_assertions(mut self, mut more: Vec<Assertion>) -> Self {
        self.assert.append(&mut more);
        self
    }
}

#[derive(Debug, Copy, Clone, PartialEq)]
pub enum Acceptance {
    Accept,
    Assertion,
    NotYet,
}

#[derive(Debug, Clone)]
pub struct EState<S: Stage> {
    pub trans: HashSet<Transition<S>>,
    pub accept: Acceptance,
}

impl EState<Epsilon> {
    fn remove_epsilon(self) -> EState<Resolved> {
        let trans = self
            .trans
            .into_iter()
            .filter_map(Transition::non_epsilon)
            .collect();
        let accept = self.accept;
        EState { trans, accept }
    }

    fn set_epsilon_transitions(&mut self, trans: impl IntoIterator<Item = Transition<Epsilon>>) {
        self.trans.retain(|t| t.consumes.is_some());
        for transition in trans.into_iter() {
            assert!(transition.consumes.is_none());
            self.trans.insert(transition);
        }
    }
}

impl<S: Stage> EState<S> {
    fn start() -> Self {
        Self {
            trans: HashSet::new(),
            accept: Acceptance::NotYet,
        }
    }
    fn terminal() -> Self {
        Self {
            trans: HashSet::new(),
            accept: Acceptance::Accept,
        }
    }

    fn remap(&mut self, mut f: impl FnMut(StateId) -> StateId) {
        self.trans = self
            .trans
            .iter()
            .cloned()
            .map(|mut t| {
                t.remap(&mut f);
                t
            })
            .collect();
    }

    fn reachable_states(&self) -> impl Iterator<Item = StateId> {
        self.trans.iter().flat_map(|t| t.reachable_states())
    }
}

#[derive(Debug)]
pub enum EnfaTranslationError {
    AssertionsNotSupported,
}

impl TryFrom<Pattern> for ENFA<Epsilon> {
    type Error = EnfaTranslationError;

    fn try_from(value: Pattern) -> Result<Self, Self::Error> {
        Ok(match value {
            Pattern::Byte(c) => Self::try_from(Pattern::Range(c, c))?,
            Pattern::Range(c1, c2) => Self {
                states: vec![
                    EState {
                        trans: set![Transition::new(ByteRange::new_range(c1, c2), 1)],
                        accept: Acceptance::NotYet,
                    },
                    EState::terminal(),
                ],
            },
            Pattern::Alt(alts) => {
                let nfas: Vec<ENFA<Epsilon>> = alts
                    .into_iter()
                    .map(Self::try_from)
                    .collect::<Result<_, _>>()?;
                let mut states = vec![EState::start()];
                let mut ends = vec![];
                for nfa in nfas.into_iter() {
                    let len = states.len();
                    states[0].trans.insert(Transition::epsilon(len));
                    states.append(&mut (nfa.shift(len)));
                    ends.push(states.len() - 1);
                }
                states.push(EState::terminal());
                for end in ends.into_iter() {
                    let last = states.len() - 1;
                    states[end].trans.insert(Transition::epsilon(last));
                }
                Self { states }
            }
            Pattern::Concat(seq) => {
                let nfas: Vec<Self> = seq
                    .into_iter()
                    .map(Self::try_from)
                    .collect::<Result<_, _>>()?;
                Self::concat(nfas)
            }
            Pattern::Rep(regex, min, None) => {
                let nfa = ENFA::try_from(*regex)?;
                let base = nfa.clone().repeat(min as usize);
                let tail = nfa.looping();
                Self::concat(vec![base, tail])
            }
            Pattern::Rep(regex, min, Some(max)) => {
                assert!(min < max);
                let nfa = Self::try_from(*regex)?;
                let base = nfa.clone().repeat(min as usize);
                let tail = nfa.optx((max - min) as usize);
                Self::concat(vec![base, tail])
            }
            Pattern::Nothing => Self {
                states: vec![EState::terminal()],
            },
            Pattern::Assertion(dir, polarity, pat) => {
                if dir == LookDirection::Behind {
                    return Err(EnfaTranslationError::AssertionsNotSupported);
                }
                let mut regex = Self::try_from(*pat)?;
                for s in regex.states.iter_mut() {
                    if s.accept == Acceptance::Accept {
                        s.accept = Acceptance::Assertion;
                    }
                }
                let mut regex = regex.shift(1);
                let mut states = Vec::with_capacity(regex.len() + 2);
                states.push(EState {
                    trans: set![Transition {
                        to: regex.len() + 1,
                        assert: vec![Assertion { to: 1, polarity }],
                        consumes: None
                    }],
                    accept: Acceptance::NotYet,
                });
                states.append(&mut regex);
                states.push(EState::terminal());
                Self { states }
            }
        })
    }
}