use crate::*;
use pyo3::prelude::*;
use rayon::prelude::*;

// Node in our character trie
#[derive(Default)]
struct CharTrieNode {
    // Store character and node pairs, keep sorted for binary search
    children: Vec<(char, CharTrieNode)>,
    // Store token IDs that are valid at this point
    token_id: Option<Py<PyAny>>,
}

impl CharTrieNode {
    // Add a child for character c, maintaining sorted order
    fn add_child(&mut self, c: char) -> &mut CharTrieNode {
        match self.children.binary_search_by_key(&c, |(ch, _)| *ch) {
            Ok(index) => &mut self.children[index].1,
            Err(index) => {
                // Insert at the right position to maintain sort order
                self.children.insert(index, (c, CharTrieNode::default()));
                &mut self.children[index].1
            }
        }
    }

    // Add a token id to this node
    fn add_token(&mut self, token_id: Py<PyAny>) {
        self.token_id = Some(token_id);
    }

    fn find_invalid_tokens(
        &self,
        py: Python<'_>,
        xml_validator: &XmlSchemaValidator,
    ) -> Vec<Py<PyAny>> {
        let res = self
            .children
            .iter()
            .map(|(c, n)| {
                let mut node_validator = xml_validator.clone();
                if node_validator.append_char(*c).is_ok() {
                    // If the node is valid, recursively check its children
                    n.find_invalid_tokens(py, &node_validator)
                } else {
                    // If the node is invalid, add it and its children to the list of invalid tokens
                    n.all_tokens(py)
                }
            })
            .flatten()
            .collect();
        res
    }

    // Parallel version of find_invalid_tokens for the first level of recursion
    fn find_invalid_tokens_parallel(
        &self,
        _py: Python<'_>,
        xml_validator: &XmlSchemaValidator,
    ) -> Vec<Py<PyAny>> {
        // We need to use the Python GIL to share PyAny objects across threads
        Python::with_gil(|_py| {
            // Convert children to a parallel iterator
            self.children
                .par_iter() // Use parallel iterator from rayon
                .map(|(c, n)| {
                    // Clone validator for this thread
                    let mut node_validator = xml_validator.clone();
                    
                    // Process this character
                    if node_validator.append_char(*c).is_ok() {
                        // For valid nodes, use the regular (non-parallel) method for deeper levels
                        Python::with_gil(|py| {
                            n.find_invalid_tokens(py, &node_validator)
                        })
                    } else {
                        // For invalid nodes, get all tokens
                        Python::with_gil(|py| {
                            n.all_tokens(py)
                        })
                    }
                })
                .flatten()
                .collect()
        })
    }

    fn all_tokens(&self, py: Python<'_>) -> Vec<Py<PyAny>> {
        let mut tokens = Vec::new();
        if let Some(token_id) = &self.token_id {
            tokens.push(token_id.clone_ref(py));
        }
        for (_, child) in &self.children {
            tokens.extend(child.all_tokens(py));
        }
        tokens
    }

    fn clone_ref(&self, py: Python<'_>) -> Self {
        Self {
            children: self
                .children
                .iter()
                .map(|(ch, node)| (*ch, node.clone_ref(py)))
                .collect(),
            token_id: self
                .token_id
                .iter()
                .map(|token_id| token_id.clone_ref(py))
                .nth(0),
        }
    }
}

// The main trie structure
#[derive(Default)]
pub struct CharTrie {
    root: CharTrieNode,
}

impl CharTrie {
    // Insert a token's character sequence into the trie
    pub fn insert(&mut self, token_text: &str, token_id: Py<PyAny>) {
        let mut current = &mut self.root;

        // Add each character in the token text
        for c in token_text.chars() {
            current = current.add_child(c);
        }

        // Mark this node as representing token_id
        current.add_token(token_id);
    }

    // Find all tokens that could be invalid given the current XML state
    pub fn find_invalid_tokens(
        &self,
        py: Python<'_>,
        xml_validator: &XmlSchemaValidator,
    ) -> Vec<Py<PyAny>> {
        // Use the parallel version starting at the root
        self.root.find_invalid_tokens_parallel(py, xml_validator)
    }

    pub fn clone_ref(&self, py: Python<'_>) -> Self {
        Self {
            root: self.root.clone_ref(py),
        }
    }
}