using Microsoft.SqlServer.TransactSql.ScriptDom;

namespace SQLLinter.Infrastructure.Rules;

public static class ParentMapBuilder
{
    public static Dictionary<TSqlFragment, TSqlFragment?> Build(TSqlFragment root)
    {
        var map = new Dictionary<TSqlFragment, TSqlFragment?>();
        Traverse(root, null, map);
        return map;
    }

    private static void Traverse(
        TSqlFragment node,
        TSqlFragment? parent,
        Dictionary<TSqlFragment, TSqlFragment?> map)
    {
        if (!map.ContainsKey(node))
            map[node] = parent;

        foreach (var child in node.GetChildren())
        {
            Traverse(child, node, map);
        }
    }
}


public static class ScriptDomExtensions
{
    public static IEnumerable<TSqlFragment> GetChildren(this TSqlFragment node)
    {
        var collector = new DirectChildrenCollector(node);
        node.Accept(collector);
        return collector.Children;
    }

    private class DirectChildrenCollector : TSqlFragmentVisitor
    {
        private readonly TSqlFragment _root;
        private bool _isRootVisited = false;

        public List<TSqlFragment> Children { get; } = new();

        public DirectChildrenCollector(TSqlFragment root)
        {
            _root = root;
        }

        public override void Visit(TSqlFragment fragment)
        {
            if (!_isRootVisited)
            {
                // Первый вызов — это сам root
                _isRootVisited = true;
            }
            else
            {
                // Все остальные вызовы — это прямые дети root
                Children.Add(fragment);

                // ВАЖНО: не спускаемся глубже
                return;
            }

            // Продолжаем обход только для root
            base.Visit(fragment);
        }
    }
}