// SPDX-FileCopyrightText: 2024 Aurélien Gâteau <mail@agateau.com>
//
// SPDX-License-Identifier: Apache-2.0

use tree_ds::prelude::AutomatedId;

use crate::parse::{RegexTree, Token};

pub fn get_token_len(tree: &RegexTree, target_node_id: &AutomatedId) -> usize {
    let target_node = tree.get_node_by_id(target_node_id).unwrap();
    let token = target_node.get_value().unwrap();
    match token {
        Token::Group { header, quantifier } => {
            let mut len = header.len() + quantifier.len() + 2;
            for node_id in target_node.get_children_ids() {
                len += get_token_len(tree, &node_id);
            }
            len
        }
        Token::Literal { text, quantifier } => text.len() + quantifier.len(),
        Token::Set { text, quantifier } => text.len() + quantifier.len() + 2,
        Token::Edge { symbol } => symbol.len(),
        Token::Alternative => 1,
    }
}

pub fn get_token_span(tree: &RegexTree, target_node_id: &AutomatedId) -> (usize, usize) {
    let root_node = tree.get_root_node().unwrap();

    let (found, start) = get_subtree_token_start(tree, &root_node.get_node_id(), target_node_id);
    if found {
        let token_len = get_token_len(tree, target_node_id);
        return (start, start + token_len);
    }
    panic!("Can't find node {} in tree", target_node_id);
}

fn get_subtree_token_start(
    tree: &RegexTree,
    root_id: &AutomatedId,
    target_node_id: &AutomatedId,
) -> (bool, usize) {
    let root_node = tree.get_node_by_id(root_id).unwrap();
    let mut start = 0;

    for node_id in root_node.get_children_ids() {
        if node_id == *target_node_id {
            return (true, start);
        }
        let token = tree.get_node_by_id(&node_id).unwrap().get_value().unwrap();
        if let Token::Group { header, quantifier } = token {
            // Pass the opening '(' and the header
            start += header.len() + 1;
            let (found, sub_start) = get_subtree_token_start(tree, &node_id, target_node_id);
            start += sub_start;
            if found {
                return (true, start);
            }
            // Pass the quantifier and the closing ')'
            start += quantifier.len() + 1;
        } else {
            start += get_token_len(tree, &node_id);
        }
    }
    (false, start)
}

#[cfg(test)]
mod tests {
    use super::*;

    use yare::parameterized;

    use crate::parse::parse_regex;

    #[parameterized(
        simple = { "abc", &[3] },
        simple_group = { "(abc)", &[5] },
        multi_groups = { "a(bb)ccc", &[1, 4, 3] },
        set = { "[abc]", &[5] },
        edges = { "^foo$", &[1, 3, 1] },
        alternative = { "aa|bbb", &[2, 1, 3]},
        nested_groups = { "a(bb(cc))dd", &[1, 8, 2]},
    )]
    fn test_get_token_len(regex: &str, expected_lens: &[usize]) {
        let tree = parse_regex(regex).unwrap();

        let root_node = tree.get_root_node().unwrap();

        let mut lens: Vec<usize> = Vec::<_>::new();
        for node_id in root_node.get_children_ids() {
            lens.push(get_token_len(&tree, &node_id));
        }
        assert_eq!(lens, expected_lens);
    }

    #[parameterized(
        simple = { "abc", &[(0,3)] },
        simple_group = { "(abc)", &[(0, 5)] },
        multi_groups = { "a(bb)ccc", &[(0, 1), (1, 5), (5, 8)] },
        alternative = { "aa|bbb", &[(0, 2), (2, 3), (3, 6)]},
    )]
    fn test_get_token_span_top_level(regex: &str, expected_spans: &[(usize, usize)]) {
        let tree = parse_regex(regex).unwrap();

        let root_node = tree.get_root_node().unwrap();

        let mut spans: Vec<(usize, usize)> = Vec::<_>::new();
        for node_id in root_node.get_children_ids() {
            spans.push(get_token_span(&tree, &node_id));
        }
        assert_eq!(spans, expected_spans);
    }

    #[test]
    fn test_get_token_span_nested() {
        let tree = parse_regex("(aa(?P<a>.))").unwrap();

        let root_node = tree.get_root_node().unwrap();
        let outer_group_id = root_node.get_children_ids()[0];
        let outer_group_node = tree.get_node_by_id(&outer_group_id).unwrap();

        let inner_group_id = outer_group_node.get_children_ids()[1];
        let inner_group_node = tree.get_node_by_id(&inner_group_id).unwrap();

        // Sanity check: does inner_group_id really point to the inner group?
        let token = inner_group_node.get_value().unwrap();
        if let Token::Group { header, .. } = token {
            assert_eq!(header, "?P<a>");
        } else {
            panic!("inner_group_node is not a group");
        }

        let (start, end) = get_token_span(&tree, &inner_group_id);
        assert_eq!((start, end), (3, 11));
    }
}
