using System;
using System.Collections;
using System.Collections.Generic;

namespace Dacryptero.Collections
{
    public class TreeCollection<T> : IReadOnlyCollection<T>
    {
        private class Node
        {
            public T? Item { get; }

            public TreeCollection<T>? Tree { get; }

            public bool HasItem { get; }

            public Node(T item)
            {
                Item = item;
                Tree = null;
                HasItem = true;
            }

            public Node(TreeCollection<T> tree)
            {
                Item = default;
                Tree = tree;
                HasItem = false;
            }
        }
    
        private readonly ReadOnlyMemory<Node> nodes;

        public int Count { get; }

        public TreeCollection()
        {
            nodes = Memory<Node>.Empty;
            Count = 0;
        }

        public TreeCollection(IReadOnlyList<T> items)
        {
            Memory<Node> nodes = new Node[items.Count];
            for (int i = 0; i < items.Count; ++i)
                nodes.Span[i] = new Node(items[i]);
            this.nodes = nodes; 
            Count = items.Count;
        }

        public TreeCollection(T item)
        {
            Memory<Node> nodes = new Node[1];
            nodes.Span[0] = new Node(item);
            this.nodes = nodes; 
            Count = 1;
        }

        public TreeCollection(ReadOnlySpan<T> items)
        {
            Memory<Node> nodes = new Node[items.Length];
            for (int i = 0; i < items.Length; ++i)
                nodes.Span[i] = new Node(items[i]);
            this.nodes = nodes; 
            Count = items.Length;
        }

        public TreeCollection(params T[] items)
            : this(items.AsSpan())
        {
        }

        public TreeCollection(IReadOnlyList<TreeCollection<T>> items)
        {
            Memory<Node> nodes = new Node[items.Count];
            Count = 0;
            for (int i = 0; i < items.Count; ++i)
            {
                nodes.Span[i] = new Node(items[i]);
                Count += items[i].Count;
            }
            this.nodes = nodes;
        }

        public TreeCollection(TreeCollection<T> item)
        {
            Memory<Node> nodes = new Node[1];
            nodes.Span[0] = new Node(item);
            this.nodes = nodes; 
            Count = item.Count;
        }

        public TreeCollection(ReadOnlySpan<TreeCollection<T>> items)
        {
            Memory<Node> nodes = new Node[items.Length];
            Count = 0;
            for (int i = 0; i < items.Length; ++i)
            {
                nodes.Span[i] = new Node(items[i]);
                Count += items[i].Count;
            }
            this.nodes = nodes; 
        }

        public TreeCollection(params TreeCollection<T>[] items)
            : this(items.AsSpan())
        {
        }

        public IEnumerator<T> GetEnumerator()
        {
            for (int i = 0; i < nodes.Length; ++i)
            {
                if (nodes.Span[i].HasItem)
                    yield return nodes.Span[i].Item!;
                else foreach (var item in nodes.Span[i].Tree!)
                    yield return item;
            }
        }

        IEnumerator IEnumerable.GetEnumerator()
            => GetEnumerator();
        
        public T[] ToArray()
            => ToMemory().ToArray();
        
        public ReadOnlyMemory<T> ToMemory()
        {
            Memory<T> result = new T[Count];
            ToMemory(result.Span);
            return result;
        }

        protected void ToMemory(Span<T> result)
        {
            int offset = 0;
            for (int i = 0; i < nodes.Length; ++i)
            {
                if (nodes.Span[i].HasItem)
                {
                    result[offset] = nodes.Span[i].Item!;
                    offset++;
                }
                else
                {
                    var end = offset + nodes.Span[i].Tree!.Count;
                    nodes.Span[i].Tree!.ToMemory(result[offset .. end]);
                    offset = end;
                }
            }
        }
    }
}