using Sawmill;

namespace InnovEnergy.Lib.SrcGen.Trees;

using T = ForwardTree<String>;

public static class PrefixTree
{
    public static Tree<String> Create(IEnumerable<String> strings, String root = "")
    {
        var children = strings.Select(s => new T(s)).ToList();
        var tree     = new T(root, children);

        var forwardTree = tree
                         .Rewrite(GroupByPrefix)
                         .Rewrite(ConcatUnaryNodes);
        
        return new Tree<String>(forwardTree);
    }

    private static T GroupByPrefix(T tree)
    {
        var newChildren = tree
                         .Children
                         .Where(s => s.Node.Length > 0)
                         .GroupBy(s => s.Node[..1], s => new T(s.Node[1..]))
                         .Select(g => GroupByPrefix(new T(g.Key, g.ToList())))
                         .ToList();

        return newChildren.Count == 0
             ? tree // leaf
             : new T(tree.Node, newChildren);
    }
    
    private static T ConcatUnaryNodes(ForwardTree<String> tree)
    {
        return tree.Children.Count == 1 && tree.Children[0] is var child
             ? new T(tree.Node + child.Node, child.Children)
             : tree;
    }
}