#pragma once

#include "Common.h"
#include "Palette.h"
#include "Lockable.h"

namespace Feather
{
	// A container that maps values from a local palette to a 
	template <typename T, uint MaxSize, uint MaxBits>
	struct PalettedContainer
	{
		//const Palette<T>& globalPalette;

        static constexpr size_t NUM_LONGS = Math::CeilDiv(MaxSize * MaxBits, bitsizeof(uint64));

        static constexpr size_t BITS_PER_CELL = bitsizeof(uint64);

        LockableArray<uint64, NUM_LONGS> data;

        // TODO: this could eventually be Palette<BlockState> to avoid double lookup
        // Palette of local indexs to indexs in the global palette
        IDMapper<int> palette;

        uint size = MaxSize;
        uint8 bits = MaxBits;
        uint64 mask = 0;
        uint valuesPerCell = 0;

    public:
        PalettedContainer(uint size, uint8 bits) : size(size), data() { SetBits(bits); }

        inline void SetBits(uint8 nbits)
        {
            bits = nbits;
            mask = (1 << bits) - 1;
            valuesPerCell = BITS_PER_CELL / bits;

            Log::Debug("PalettedContainer::SetBits: {}, mask = {:#b}, valuesPerCell = {}", bits, mask, valuesPerCell);
        }
        
        // TODO: optimize all math here

        inline int CellIndex(int index)
        {
            return index / valuesPerCell;
        }

        // Update value, returns previous value
        // TODO: change from uint to const T&
        inline uint GetAndSet(int index, uint value)
        {
            auto [d, lock] = data.borrow();
            
            int id = value; // palette.GetID(value);
            int cellIndex = CellIndex(index);
            uint64 shift = (index - cellIndex * valuesPerCell) * bits;//(index % valuesPerCell) * bits;

            Log::Debug("PalettedContainer::GetAndSet: index {} -> cell {}, [{}:{}]", index, cellIndex, shift, shift + bits);

            uint64 cell = d[cellIndex];
            int prev = (cell >> shift) & mask;

            //d[cellIndex] = ((cell & (mask << shift ^ 0xFFFFFFFFFFFFFFFFL)) | (uint64(id) & mask)) << shift;

            d[cellIndex] &= ~(mask << shift);               // clear target bits
            d[cellIndex] |= ((uint64)id & mask) << shift;   // set target bits

            //Log::Debug("Changed from {:#b} to {:#b}", cell, d[cellIndex]);
            Log::Debug("Changed from {} (g: {}) to {} (g: {})", prev, palette.ByID(prev), value, palette.ByID(value));

            return (uint)prev; // palette.ByID(prev);
        }

        inline constexpr size_t GetNumLongs() { return NUM_LONGS; }
	};
}