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

namespace Qrakhen.Qamp.Core.Collections;

/// <summary>
///   Similar to a register with <see cref="Address"/> as key; 
///   But with the caveat that the memory slots can be allocated or freed up, 
///   automatically seeking for free addresses when adding data,
///   making it much easier to deal with addresses or pointers.
/// </summary>
public class Table<TValue> :
   IGetSet<Address, TValue>,
   IToArray<TValue>,
   IEnumerable<KeyValuePair<Address, TValue>>
{
   public const int BLOCK_SIZE = sizeof(int) * 8;
   private const uint BLOCK_FULL = (uint)((1UL << 32) - 1);

   private readonly TValue[] _data;
   private readonly uint[] _alloc;

   public int Size { get; }
   public int Blocks => Size / BLOCK_SIZE;

   public TValue this[Address index] 
   { 
      get => Get(index);
      set => Set(index, value); 
   }

   public Table(int size = 4096)
   {
      Size = size;
      if (size % BLOCK_SIZE != 0)
         throw new ArgumentException($"{nameof(Table<TValue>)} size parameter must be product (multiple) of its block size {BLOCK_SIZE}.");
      _data = new TValue[size];
      _alloc = new uint[size / BLOCK_SIZE];
   }

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

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

   public bool IsAllocated(Address address)
   {
      return (_alloc[GetBlockIndex(address)] & (1 << GetBitIndex(address))) > 0;
   }

   private void SetAllocBit(Address address, bool value)
   {
      int block = GetBlockIndex(address);
      int index = GetBitIndex(address);
      if (value)
         _alloc[block] |= 1U << index;
      else
         _alloc[block] &= ~(1U << index);
   }

   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;
   }

   private bool Validate(Address address)
   {
      if (address < 0 || address >= Size)
         throw new RuntimeException($"Tried to access address outside memory {address}");
      return true;
   }

   public void Free(Address address, int size = 1)
   {
      ArgumentOutOfRangeException.ThrowIfNotEqual(1, size); // only support single address freeing atm
      SetAllocBit(address, false);
   }

   public Address Allocate(int size = 1)
   {
      ArgumentOutOfRangeException.ThrowIfNotEqual(1, size); // only support single address allocation atm
      Address address = SeekFree();
      SetAllocBit(address, true);
      return address;
   }

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

   public TValue Get(Address index)
   {
      Validate(index);
      if (!IsAllocated(index))
         return default;
      return _data[index];
   }

   public bool TryGet(Address index, [MaybeNullWhen(false)] out TValue? value)
   {
      Validate(index);
      value = default;
      if (!IsAllocated(index))
         return false;
      value = _data[index];
      return true;
   }

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

   /// <summary>
   ///   Calls <see cref="Allocate(int)"/> and then <see cref="Set(Address, TValue)"/>, 
   ///   returning the address at which <paramref name="value"/> was stored at.
   /// </summary>
   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 string Print()
   {
      string r = $"Table [{nameof(TValue)}]:\n";
      int c = 0;
      for (int b = 0; b < Blocks; b++)
      {
         uint block = _alloc[b];
         if (block == 0)
            continue;
         for (int i = 0; i < BLOCK_SIZE; i++)
         {
            Address address = MakeAddress(b, i);
            if (!IsAllocated(address))
               continue;
            c++;
            r += $"  {address}: {Get(address)}\n";
         }
      }
      r += $"{c}/{Size - c} slots allocated.";
      return r;
   }

   public static Address MakeAddress(int block, int index) => new((long)block * BLOCK_SIZE + index);

   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() => GetEnumerator();
}