using Qrakhen.Qamp.Core.Collections.Abstractions;
using Qrakhen.Qamp.Core.Values;
using System.Collections;

namespace Qrakhen.Qamp.Core.Collections;

/// <summary>
///   Similar to a register with the caveat that the memory can free up or reserve blocks within itself, returning the next writeable address when storing data.
/// </summary>
public class Memory<TValue> :
   IGetSet<Address, TValue>,
   IToArray<TValue>,
   IEnumerable<KeyValuePair<Address, TValue>>
{
   private const int BLOCK_SIZE = sizeof(int);
   private const uint BLOCK_FULL = (uint)((1UL << 32) - 1);

   private readonly TValue[] _data;

   private readonly uint[] _alloc;

   public int Size { get; }

   public TValue this[Address index] 
   { 
      get => _data[index]; 
      set => _data[index] = value; 
   }

   public Memory(int size = 4096)
   {
      Size = size;
      _data = new TValue[size];
   }

   private int GetBlockIndex(Address address)
   {
      return (int)(address / BLOCK_SIZE);
   }

   private int GetBitIndex(Address address)
   {
      return (int)(address % BLOCK_SIZE);
   }

   private Address SeekFree(Address from = default)
   {
      Address cur = from - (from % BLOCK_SIZE);
      while (cur < Size)
      {
         int block = GetBlockIndex(cur);
         uint mask = _alloc[block];
         if (mask < BLOCK_FULL)
         {
            int index = 0;
            while ((mask & (1 << index)) > 0)            
               index++;

            return (Address)((long)block * BLOCK_SIZE + index);
         }
         cur += BLOCK_SIZE;
      }
      return Address.Void;
   }

   public void Free(Address address, int size = 1)
   {
      ArgumentOutOfRangeException.ThrowIfNotEqual(1, size); // only support single address freeing atm
      int block = GetBlockIndex(address);
      int index = GetBitIndex(address);
      uint flags = _alloc[block];
      _alloc[block] &= ~(1U << index);
   }

   public void Allocate(int size = 1)
   {
      ArgumentOutOfRangeException.ThrowIfNotEqual(1, size); // only support single address allocation atm
      Address address = SeekFree();
      int block = GetBlockIndex(address);
      int index = GetBitIndex(address);
      _alloc[block] |= (1U << index);
   }

   public TValue[] ToArray()
   {
      return _data.ToArray();
   }

   public TValue Get(Address index)
   {
      return _data[index];
   }

   public void Set(Address index, TValue value)
   {
      _data[index] = value;
   }

   public Address Add(TValue value)
   {
      Address free = SeekFree();
      if (free == Address.Void) 
         throw new QampException($"Tried to add {value} to a memory block, but no free memory could be allocated.");
      Set(free, value);
      return free;
   }

   public IEnumerator<KeyValuePair<Address, TValue>> GetEnumerator()
   {
      for (Address addr = 0; addr < Size; addr++)
         yield return new KeyValuePair<Address, TValue>(addr, Get(addr));
   }

   IEnumerator IEnumerable.GetEnumerator()
   {
      return GetEnumerator();
   }
}