//! Build a decision tree from CSS selectors.
//!
//! The tree weaves selectors so shared conditions are checked once. The build
//! pipeline parses each selector, rewrites class and id selectors as attribute
//! selectors, normalizes, then weaves the results into branches.

use std::borrow::Cow;
use std::collections::HashSet;
use std::hash::Hash;
use std::sync::Arc;

use crate::ast::{AttrPredicate, DecisionTreeNode, MatcherNode, ValueContainer, VariantNode};
use crate::css::ast::{CompoundSelector, Matcher, Modifier, SimpleSelector};
use crate::css::{normalize, parse1, serialize_compound, NormalizeOptions};
use crate::error::SelectorError;
use crate::ordered_map::OrderedMap;

/// Options for [`DecisionTree`] construction.
#[derive(Debug, Clone, Default)]
pub struct DecisionTreeOptions {
    /// Attribute names whose values default to case-insensitive.
    ///
    /// Only affects attribute value selectors with no explicit `i` or `s`
    /// modifier.
    pub attributes_with_normalized_values: Vec<String>,
}

/// A CSS selectors decision tree.
///
/// Weaves similar selectors together so an element needs the fewest checks to
/// find every matching selector. Turn it into a callable form with a builder
/// via [`DecisionTree::build`].
///
/// `V` is the type of values associated with selectors.
pub struct DecisionTree<V> {
    branches: Vec<DecisionTreeNode<V>>,
}

impl<V> DecisionTree<V> {
    /// Build a decision tree from selector and value pairs.
    ///
    /// Invalid input is a routine, caller-controlled outcome, so this returns a
    /// [`SelectorError`] rather than panicking.
    ///
    /// # Errors
    ///
    /// Returns [`SelectorError::UnsupportedKind`] for an unsupported combinator
    /// or functional pseudo-class, [`SelectorError::UnsupportedPseudoClass`] for
    /// a pseudo-class outside the supported set,
    /// [`SelectorError::UnsupportedNamespace`] for a namespace prefix, and
    /// [`SelectorError::Parse`] for a tokenizer or parser failure.
    pub fn new(
        input: Vec<(String, V)>,
        options: DecisionTreeOptions,
    ) -> Result<Self, SelectorError> {
        let pairs = to_ast_terminal_pairs(input, &options)?;
        let branches = weave(pairs)?;
        Ok(DecisionTree { branches })
    }

    /// Turn this tree into a callable form using `builder`.
    ///
    /// `builder` receives the root node slice and returns whatever it produces.
    /// The built-in [`crate::treeify`] builder returns a `String`. A DOM plugin
    /// returns a [`crate::Picker`].
    pub fn build<R>(&self, builder: impl FnOnce(&[DecisionTreeNode<V>]) -> R) -> R {
        builder(&self.branches)
    }
}

// ==============================================

/// The terminal carried alongside an in-progress selector AST.
enum Terminal<V> {
    /// Records a final match.
    Match(ValueContainer<V>),
    /// Returns to the element beneath after the left side is exhausted.
    PopElement(Vec<DecisionTreeNode<V>>),
}

/// A selector AST paired with the terminal to emit once it is consumed.
struct AstTerminalPair<V> {
    ast: CompoundSelector,
    terminal: Terminal<V>,
}

/// One group of items sharing a key, with a representative simple selector.
struct SelectorsGroup<V> {
    one_simple_selector: SimpleSelector,
    items: Vec<AstTerminalPair<V>>,
}

/// The three groups produced by [`break_by_kind`]: matches, non-matches, and
/// exhausted items.
type BrokenItems<V> = (
    Vec<AstTerminalPair<V>>,
    Vec<AstTerminalPair<V>>,
    Vec<AstTerminalPair<V>>,
);

/// Parse, reduce, and normalize each input selector into an item.
fn to_ast_terminal_pairs<V>(
    input: Vec<(String, V)>,
    options: &DecisionTreeOptions,
) -> Result<Vec<AstTerminalPair<V>>, SelectorError> {
    let normalized_attrs: HashSet<String> = options
        .attributes_with_normalized_values
        .iter()
        .map(|s| s.to_lowercase())
        .collect();
    let norm_opts = NormalizeOptions {
        attributes_with_normalized_values: normalized_attrs,
    };

    let mut results = Vec::with_capacity(input.len());
    for (index, (selector_string, value)) in input.into_iter().enumerate() {
        let mut ast = parse1(&selector_string)?;
        reject_namespaces(&ast)?;
        reduce_selector_variants(&mut ast);
        normalize(&mut ast, &norm_opts);
        let specificity = ast.specificity;
        results.push(AstTerminalPair {
            ast,
            terminal: Terminal::Match(ValueContainer {
                index,
                specificity,
                value,
            }),
        });
    }
    Ok(results)
}

/// Rewrite class and id selectors as attribute selectors and drop universals.
///
/// Recurses into combinator left sides. The compound-level specificity stays
/// as parsed, so an id keeps `[1, 0, 0]` even after the rewrite.
fn reduce_selector_variants(ast: &mut CompoundSelector) {
    let mut new_list = Vec::with_capacity(ast.list.len());
    for sel in ast.list.drain(..) {
        match sel {
            SimpleSelector::Class {
                name, specificity, ..
            } => {
                new_list.push(SimpleSelector::AttrValue {
                    name: "class".to_string(),
                    namespace: None,
                    matcher: Matcher::Includes,
                    value: name,
                    modifier: None,
                    specificity,
                });
            }
            SimpleSelector::Id { name, .. } => {
                new_list.push(SimpleSelector::AttrValue {
                    name: "id".to_string(),
                    namespace: None,
                    matcher: Matcher::Equals,
                    value: name,
                    modifier: None,
                    // Safe change: simple-selector specificity is not read after this.
                    specificity: [0, 1, 0],
                });
            }
            SimpleSelector::Combinator {
                combinator,
                mut left,
                specificity,
            } => {
                reduce_selector_variants(&mut left);
                new_list.push(SimpleSelector::Combinator {
                    combinator,
                    left,
                    specificity,
                });
            }
            SimpleSelector::Universal { .. } => {
                // skip it
            }
            other => new_list.push(other),
        }
    }
    ast.list = new_list;
}

/// Reject any namespace prefix.
///
/// The tree branches on the local name and would drop the prefix, widening the
/// match. Reject at build time instead, the way unsupported combinators are.
fn reject_namespaces(ast: &CompoundSelector) -> Result<(), SelectorError> {
    for sel in &ast.list {
        let prefix = match sel {
            SimpleSelector::Universal { namespace, .. }
            | SimpleSelector::Tag { namespace, .. }
            | SimpleSelector::AttrPresence { namespace, .. }
            | SimpleSelector::AttrValue { namespace, .. } => namespace.as_deref(),
            SimpleSelector::Combinator { left, .. } => {
                reject_namespaces(left)?;
                None
            }
            _ => None,
        };
        if let Some(prefix) = prefix {
            return Err(SelectorError::UnsupportedNamespace {
                prefix: prefix.to_string(),
            });
        }
    }
    Ok(())
}

/// Weave items into a list of sibling branches.
///
/// Each pass picks the most frequent selector kind, splits matching items into
/// one branch, terminates exhausted items, and continues with the rest.
fn weave<V>(mut items: Vec<AstTerminalPair<V>>) -> Result<Vec<DecisionTreeNode<V>>, SelectorError> {
    let mut branches = Vec::new();
    while !items.is_empty() {
        let top_kind = find_top_key(&items, |_| true, get_selector_kind);
        let (matches, non_matches, empty) = break_by_kind(items, top_kind.as_ref());
        items = non_matches;
        if !matches.is_empty() {
            branches.push(branch_of_kind(top_kind.as_ref(), matches)?);
        }
        if !empty.is_empty() {
            branches.extend(terminate(empty));
        }
    }
    Ok(branches)
}

/// Emit terminals for exhausted items. Lift direct terminals out of popElement.
fn terminate<V>(items: Vec<AstTerminalPair<V>>) -> Vec<DecisionTreeNode<V>> {
    let mut results = Vec::new();
    for item in items {
        match item.terminal {
            Terminal::Match(vc) => {
                results.push(DecisionTreeNode::Terminal {
                    value_container: vc,
                });
            }
            Terminal::PopElement(cont) => {
                let mut rest = Vec::new();
                for node in cont {
                    if matches!(node, DecisionTreeNode::Terminal { .. }) {
                        results.push(node);
                    } else {
                        rest.push(node);
                    }
                }
                if !rest.is_empty() {
                    results.push(DecisionTreeNode::PopElement { cont: rest });
                }
            }
        }
    }
    results
}

/// Split items into matches, non-matches, and exhausted items.
///
/// Items are not modified. An item with an empty list is exhausted.
fn break_by_kind<V>(items: Vec<AstTerminalPair<V>>, selected_kind: &str) -> BrokenItems<V> {
    let mut matches = Vec::new();
    let mut non_matches = Vec::new();
    let mut empty = Vec::new();
    for item in items {
        if item.ast.list.is_empty() {
            empty.push(item);
        } else if item
            .ast
            .list
            .iter()
            .any(|node| get_selector_kind(node).as_ref() == selected_kind)
        {
            matches.push(item);
        } else {
            non_matches.push(item);
        }
    }
    (matches, non_matches, empty)
}

/// The kind key for a simple selector, used to group items into branches.
///
/// Fixed kinds borrow a static literal. Named kinds allocate the `name`-bearing
/// key. `branch_of_kind` parses these keys by prefix, so the text must stay
/// byte-identical.
fn get_selector_kind(sel: &SimpleSelector) -> Cow<'static, str> {
    match sel {
        SimpleSelector::AttrPresence { name, .. } => Cow::Owned(format!("attrPresence {}", name)),
        SimpleSelector::AttrValue { name, .. } => Cow::Owned(format!("attrValue {}", name)),
        SimpleSelector::PseudoClass { name, .. } => Cow::Owned(format!("pc {}", name)),
        SimpleSelector::Combinator { combinator, .. } => {
            Cow::Owned(format!("combinator {}", combinator))
        }
        SimpleSelector::Tag { .. } => Cow::Borrowed("tag"),
        SimpleSelector::Universal { .. } => Cow::Borrowed("universal"),
        SimpleSelector::Class { .. } => Cow::Borrowed("class"),
        SimpleSelector::Id { .. } => Cow::Borrowed("id"),
        SimpleSelector::FunctionalPseudoClass { name, .. } => Cow::Owned(format!("fpc:{}", name)),
    }
}

/// Build the branch for the given kind and matched items.
fn branch_of_kind<V>(
    kind: &str,
    items: Vec<AstTerminalPair<V>>,
) -> Result<DecisionTreeNode<V>, SelectorError> {
    if kind == "tag" {
        return tag_name_branch(items);
    }
    if let Some(name) = kind.strip_prefix("attrValue ") {
        return attr_value_branch(name, items);
    }
    if let Some(name) = kind.strip_prefix("attrPresence ") {
        return attr_presence_branch(name, items);
    }
    if let Some(name) = kind.strip_prefix("pc ") {
        return pseudo_class_branch(name, items);
    }
    if kind == "combinator >" {
        return combinator_branch('>', items);
    }
    if kind == "combinator +" {
        return combinator_branch('+', items);
    }
    Err(SelectorError::UnsupportedKind {
        kind: kind.to_string(),
    })
}

fn tag_name_branch<V>(
    items: Vec<AstTerminalPair<V>>,
) -> Result<DecisionTreeNode<V>, SelectorError> {
    let groups = splice_and_group(
        items,
        |x| matches!(x, SimpleSelector::Tag { .. }),
        |x| match x {
            SimpleSelector::Tag { name, .. } => name.clone(),
            _ => unreachable!(),
        },
    );
    let mut variants = Vec::new();
    for (name, group) in groups.into_iter() {
        variants.push(VariantNode {
            value: name,
            cont: weave(group.items)?,
        });
    }
    Ok(DecisionTreeNode::TagName { variants })
}

fn attr_presence_branch<V>(
    name: &str,
    mut items: Vec<AstTerminalPair<V>>,
) -> Result<DecisionTreeNode<V>, SelectorError> {
    for item in &mut items {
        let spliced = splice_simple_selector(
            item,
            |x| matches!(x, SimpleSelector::AttrPresence { name: n, .. } if n == name),
        );
        // branch_of_kind selected this kind from items that contain it, so
        // every item here has the selector. The splice always succeeds.
        debug_assert!(spliced.is_some(), "attrPresence {name} not found in item");
    }
    Ok(DecisionTreeNode::AttrPresence {
        name: name.to_string(),
        cont: weave(items)?,
    })
}

fn attr_value_branch<V>(
    name: &str,
    items: Vec<AstTerminalPair<V>>,
) -> Result<DecisionTreeNode<V>, SelectorError> {
    let groups = splice_and_group(
        items,
        |x| matches!(x, SimpleSelector::AttrValue { name: n, .. } if n == name),
        |x| match x {
            SimpleSelector::AttrValue {
                matcher,
                modifier,
                value,
                ..
            } => format!(
                "{} {} {}",
                matcher.as_str(),
                modifier.map(|m| m.as_str()).unwrap_or(""),
                value
            ),
            _ => unreachable!(),
        },
    );
    let mut matchers = Vec::new();
    for (_key, group) in groups.into_iter() {
        if let SimpleSelector::AttrValue {
            matcher,
            modifier,
            value,
            ..
        } = &group.one_simple_selector
        {
            matchers.push(MatcherNode {
                matcher: *matcher,
                modifier: *modifier,
                value: value.clone(),
                predicate: get_attr_value_predicate(*matcher, *modifier, value),
                cont: weave(group.items)?,
            });
        }
    }
    Ok(DecisionTreeNode::AttrValue {
        name: name.to_string(),
        matchers,
    })
}

/// ASCII whitespace, the separators for `~=` token lists.
const ASCII_WHITESPACE: [char; 5] = [' ', '\t', '\n', '\x0c', '\r'];

/// Build the predicate closure for an attribute value matcher.
fn get_attr_value_predicate(
    matcher: Matcher,
    modifier: Option<Modifier>,
    value: &str,
) -> AttrPredicate {
    let case_insensitive = modifier == Some(Modifier::CaseInsensitive);
    let expected: String = if case_insensitive {
        value.to_lowercase()
    } else {
        value.to_string()
    };
    let test: Arc<dyn Fn(&str) -> bool + Send + Sync> = if case_insensitive {
        match matcher {
            Matcher::Equals => Arc::new(move |actual: &str| expected == actual.to_lowercase()),
            Matcher::Includes => Arc::new(move |actual: &str| {
                actual
                    .to_lowercase()
                    .split(ASCII_WHITESPACE)
                    .any(|tok| tok == expected)
            }),
            Matcher::Prefix => {
                Arc::new(move |actual: &str| actual.to_lowercase().starts_with(&expected))
            }
            Matcher::Suffix => {
                Arc::new(move |actual: &str| actual.to_lowercase().ends_with(&expected))
            }
            Matcher::Substring => {
                Arc::new(move |actual: &str| actual.to_lowercase().contains(&expected))
            }
            Matcher::DashMatch => {
                Arc::new(move |actual: &str| dash_match(&actual.to_lowercase(), &expected))
            }
        }
    } else {
        match matcher {
            Matcher::Equals => Arc::new(move |actual: &str| expected == actual),
            Matcher::Includes => Arc::new(move |actual: &str| {
                actual.split(ASCII_WHITESPACE).any(|tok| tok == expected)
            }),
            Matcher::Prefix => Arc::new(move |actual: &str| actual.starts_with(&expected)),
            Matcher::Suffix => Arc::new(move |actual: &str| actual.ends_with(&expected)),
            Matcher::Substring => Arc::new(move |actual: &str| actual.contains(&expected)),
            Matcher::DashMatch => Arc::new(move |actual: &str| dash_match(actual, &expected)),
        }
    };
    AttrPredicate {
        matcher,
        modifier,
        value: value.to_string(),
        test,
    }
}

/// `|=` rule: exact match or `expected-` prefix.
fn dash_match(actual: &str, expected: &str) -> bool {
    if actual == expected {
        return true;
    }
    if let Some(rest) = actual.strip_prefix(expected) {
        return rest.starts_with('-');
    }
    false
}

fn pseudo_class_branch<V>(
    name: &str,
    mut items: Vec<AstTerminalPair<V>>,
) -> Result<DecisionTreeNode<V>, SelectorError> {
    if name != "empty"
        && name != "only-child"
        && name != "first-child"
        && name != "last-child"
        && name != "any-link"
    {
        return Err(SelectorError::UnsupportedPseudoClass {
            name: name.to_string(),
        });
    }
    for item in &mut items {
        let spliced = splice_simple_selector(
            item,
            |x| matches!(x, SimpleSelector::PseudoClass { name: n, .. } if n == name),
        );
        // The kind was pre-selected for these items, so the splice always finds
        // the pseudo-class.
        debug_assert!(spliced.is_some(), "pc {name} not found in item");
    }
    Ok(DecisionTreeNode::PseudoClass {
        name: name.to_string(),
        cont: weave(items)?,
    })
}

fn combinator_branch<V>(
    combinator: char,
    items: Vec<AstTerminalPair<V>>,
) -> Result<DecisionTreeNode<V>, SelectorError> {
    let comb_str = combinator.to_string();
    let groups = splice_and_group(
        items,
        |x| matches!(x, SimpleSelector::Combinator { combinator: c, .. } if *c == comb_str),
        |x| match x {
            SimpleSelector::Combinator { left, .. } => serialize_compound(left),
            _ => unreachable!(),
        },
    );
    let mut left_items = Vec::new();
    for (_key, group) in groups.into_iter() {
        let right_cont = weave(group.items)?;
        if let SimpleSelector::Combinator { left, .. } = group.one_simple_selector {
            left_items.push(AstTerminalPair {
                ast: *left,
                terminal: Terminal::PopElement(right_cont),
            });
        }
    }
    Ok(DecisionTreeNode::PushElement {
        combinator,
        cont: weave(left_items)?,
    })
}

/// Repeatedly pull the most common key out of the items, grouping by key.
///
/// Each pass partitions items into those carrying the current top key and the
/// rest, then splices that key out of the matched items. Every matched item
/// holds the key, so the splice that yields the representative selector always
/// succeeds.
fn splice_and_group<V, P, K>(
    mut items: Vec<AstTerminalPair<V>>,
    predicate: P,
    key_callback: K,
) -> OrderedMap<String, SelectorsGroup<V>>
where
    P: Fn(&SimpleSelector) -> bool,
    K: Fn(&SimpleSelector) -> String,
{
    let mut groups: OrderedMap<String, SelectorsGroup<V>> = OrderedMap::new();
    while !items.is_empty() {
        let best_key = find_top_key(&items, |s| predicate(s), |s| key_callback(s));
        let best_key_predicate =
            |sel: &SimpleSelector| predicate(sel) && key_callback(sel) == best_key;

        let mut matches = Vec::new();
        let mut rest = Vec::new();
        for item in items {
            if item.ast.list.iter().any(&best_key_predicate) {
                matches.push(item);
            } else {
                rest.push(item);
            }
        }

        let mut one_simple_selector: Option<SimpleSelector> = None;
        for item in &mut matches {
            let spliced = splice_simple_selector(item, best_key_predicate);
            if one_simple_selector.is_none() {
                one_simple_selector = spliced;
            }
        }
        let one = one_simple_selector
            .expect("every matched item carries the best key, so a selector is always spliced");
        groups.insert(
            best_key,
            SelectorsGroup {
                one_simple_selector: one,
                items: matches,
            },
        );
        items = rest;
    }
    groups
}

/// Remove every matching simple selector from one compound, return the first.
fn splice_simple_selector<V, P>(
    item: &mut AstTerminalPair<V>,
    predicate: P,
) -> Option<SimpleSelector>
where
    P: Fn(&SimpleSelector) -> bool,
{
    let list = &mut item.ast.list;
    let mut first_index: Option<usize> = None;
    // Scan downward so the lowest matching index wins as the representative.
    for i in (0..list.len()).rev() {
        if predicate(&list[i]) {
            first_index = Some(i);
        }
    }
    let first_index = first_index?;
    let result = list[first_index].clone();
    let mut kept = Vec::with_capacity(list.len());
    for sel in list.drain(..) {
        if !predicate(&sel) {
            kept.push(sel);
        }
    }
    item.ast.list = kept;
    Some(result)
}

/// Find the key that appears in the most items.
///
/// Each key is counted once per item. Ties go to the first key to reach the
/// running maximum, in first-encountered order. Returns the default key (empty
/// string for `String` keys) when no item carries a matching selector.
fn find_top_key<V, P, K, S>(items: &[AstTerminalPair<V>], predicate: P, key_callback: K) -> S
where
    P: Fn(&SimpleSelector) -> bool,
    K: Fn(&SimpleSelector) -> S,
    S: Clone + Eq + Hash + Default,
{
    let mut candidates: OrderedMap<S, u32> = OrderedMap::new();
    for item in items {
        let mut seen: OrderedMap<S, bool> = OrderedMap::new();
        for node in item.ast.list.iter().filter(|s| predicate(s)) {
            seen.insert(key_callback(node), true);
        }
        for key in seen.keys() {
            let entry = candidates.get_mut_or_insert(key.clone(), 0);
            *entry += 1;
        }
    }
    let mut top_kind = S::default();
    let mut top_counter = 0;
    for (key, count) in candidates.iter() {
        if *count > top_counter {
            top_kind = key.clone();
            top_counter = *count;
        }
    }
    top_kind
}
