using System;
using System.Collections.Generic;
using System.Linq;
using Avalonia.Controls;
using Avalonia.Styling;
using Avalonia.VisualTree;

namespace Avalonia.Controls.Templates
{
    public static class TemplateExtensions
    {
        public static IEnumerable<IControl> GetTemplateChildren(this ITemplatedControl control)
        {
            foreach (IControl child in GetTemplateChildren((IControl)control, control))
            {
                yield return child;
            }
        }

        private static IEnumerable<IControl> GetTemplateChildren(IControl control, ITemplatedControl templatedParent)
        {
            foreach (IControl child in control.GetVisualChildren())
            {
                var childTemplatedParent = child.TemplatedParent;

                if (childTemplatedParent == templatedParent)
                {
                    yield return child;
                }

                if (childTemplatedParent != null)
                {
                    foreach (var descendant in GetTemplateChildren(child, templatedParent))
                    {
                        yield return descendant;
                    }
                }
            }
        }
    }
}
