use biome_rowan::{
    AstNode, AstSeparatedElement, AstSeparatedList, Language, SyntaxError, SyntaxNode, SyntaxToken,
    chain_trivia_pieces, trim_trailing_trivia_pieces,
};
use std::cmp::Ordering;

/// Returns `true` if `list` is sorted by `get_key`.
/// The function returns an error if we encounter a buggy node or separator.
///
/// The list is divided into chunks of nodes with keys.
/// Thus, a node without key acts as a chuck delimiter.
/// Chunks are sorted separately.
pub fn is_separated_list_sorted_by<'a, L: Language + 'a, N: AstNode<Language = L> + 'a, Key>(
    list: &impl AstSeparatedList<Language = L, Node = N>,
    get_key: impl Fn(&N) -> Option<Key>,
    comparator: impl Fn(&Key, &Key) -> Ordering,
) -> Result<bool, SyntaxError> {
    let mut is_sorted = true;

    if list.len() > 1 {
        let mut previous_key: Option<Key> = None;
        for AstSeparatedElement {
            node,
            trailing_separator,
        } in list.elements()
        {
            // We have to check if the separator is not buggy.
            let _separator = trailing_separator?;
            previous_key = if let Some(key) = get_key(&node?) {
                if previous_key.is_some_and(|previous_key| comparator(&previous_key, &key).is_gt())
                {
                    // We don't return early because we want to return the error if we met one.
                    is_sorted = false;
                }
                Some(key)
            } else {
                // If a name cannot be extracted, then the current chunk stops here.
                None
            };
        }
    }
    Ok(is_sorted)
}

/// Returns the items and their separators resulting from sorting `list` by `get_key`.
/// When elements are reordered, `make_separator` is called to add missing separators in the middle of the list.
///
/// The list is divided into chunks of nodes with keys.
/// Thus, a node without key acts as a chuck delimiter.
/// Chunks are sorted separately.
///
/// This sort is stable (i.e., does not reorder equal elements).
pub fn sorted_separated_list_by<'a, L: Language + 'a, List, Node, Key>(
    list: &List,
    get_key: impl Fn(&Node) -> Option<Key>,
    make_separator: fn() -> SyntaxToken<L>,
    comparator: impl Fn(&Key, &Key) -> Ordering,
) -> Result<List, SyntaxError>
where
    List: AstSeparatedList<Language = L, Node = Node> + AstNode<Language = L> + 'a,
    Node: AstNode<Language = L> + 'a,
{
    let mut elements = Vec::with_capacity(list.len());
    for AstSeparatedElement {
        node,
        trailing_separator,
    } in list.elements()
    {
        let node = node?;
        let trailing_separator = trailing_separator?;
        elements.push((get_key(&node), node, trailing_separator));
    }

    // Iterate over chunks of node with a key
    for slice in elements.split_mut(|(key, _, _)| key.is_none()) {
        let last_has_separator = slice.last().is_some_and(|(_, _, sep)| sep.is_some());
        slice.sort_by(|(key1, _, _), (key2, _, _)| match (key1, key2) {
            (Some(k1), Some(k2)) => comparator(k1, k2),
            (Some(_), None) => Ordering::Greater,
            (None, Some(_)) => Ordering::Less,
            (None, None) => Ordering::Equal,
        });
        fix_separators(
            slice.iter_mut().map(|(_, node, sep)| (node, sep)),
            last_has_separator,
            make_separator,
        );
    }

    let separators: Vec<_> = elements
        .iter_mut()
        .filter_map(|(_, _, sep)| sep.take())
        .collect();
    let mut separators = separators.into_iter();
    let mut items = elements.into_iter().map(|(_, node, _)| node);

    Ok(List::unwrap_cast(SyntaxNode::new_detached(
        list.syntax().kind(),
        (0..list.len() + separators.len()).map(|index| {
            if index % 2 == 0 {
                Some(items.next()?.into_syntax().into())
            } else {
                Some(separators.next()?.into())
            }
        }),
    )))
}

/// Fix the ordered sequence of nodes and separators adding missing separators and removing an extra separator.
///
/// If a separator is missing in the middle of the sequence, then a new one is created using `make_separator`.
/// If the last node has no separator, then a new one is created only if `needs_last_separator` is set to `true`.
/// If the last node has a separator and `needs_last_separator` is set to false, then the separator is removed.
/// The separator is always kept if some comments are attached.
///
/// This utility is notably useful when a delimited list with an optional last separator is reordered.
/// It allows to add missing separators and remove an extra separator.
/// Usually, you collect every pair of nodes and separators in a vector and then pass a mutable iterator to `fix_separators`.
///
/// See [sorted_separated_list_by] as a usage example.
pub fn fix_separators<'a, L: Language + 'a, N: AstNode<Language = L> + 'a>(
    // Mutable iterator of a list of nodes and their optional separators
    iter: impl std::iter::ExactSizeIterator<Item = (&'a mut N, &'a mut Option<SyntaxToken<L>>)>,
    needs_last_separator: bool,
    make_separator: fn() -> SyntaxToken<L>,
) {
    let last_index = iter.len().saturating_sub(1);
    for (i, (node, optional_separator)) in iter.enumerate() {
        if let Some(separator) = optional_separator {
            // Remove the last separator at the separator has no attached comments
            if i == last_index
                && !(needs_last_separator
                    || separator.has_leading_comments()
                    || separator.has_trailing_comments())
            {
                // Transfer the separator trivia
                if let Some(new_node) = node.clone().append_trivia_pieces(chain_trivia_pieces(
                    separator.leading_trivia().pieces(),
                    separator.trailing_trivia().pieces(),
                )) {
                    *node = new_node;
                }
                *optional_separator = None;
            }
        } else if i != last_index || needs_last_separator {
            // The last node is moved and has no trailing separator.
            // Thus we build a new separator and remove its trailing trivia.
            *optional_separator = Some(match node.syntax().last_trailing_trivia() {
                // Transfer the trailing trivia to the separator
                Some(trivia) => make_separator()
                    .append_trivia_pieces(trim_trailing_trivia_pieces(trivia.pieces())),
                _ => make_separator(),
            });

            if let Some(new_node) = node.clone().with_trailing_trivia_pieces([]) {
                *node = new_node;
            }
        }
    }
}
