#pragma once

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

namespace Feather
{
	// A container that maps values from a local palette to a global palette
    // Implemented like Minecraft's BitStorage in 1.15.2 (with cell overflow)
	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;

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

        inline void SetBits(uint8 nbits)
        {
            bits = nbits;
            mask = (1UL << bits) - 1UL;
        }
        
        // TODO: optimize all math here

        // Update value, returns previous value
        // TODO: change from uint to const T&
        inline uint GetAndSet(int index, uint value)
        {
            // Note: Java >> is signed, while >>> is used for zero-fill shift.
            // Usage of signed >> instead of >>> has been marked.

            // Bit # we want to read from
            int bitIndex = index * bits;

            // Cell index
            uint cell = bitIndex >> 6;  // java signed >>
            // Next cell, if bits overflow
            uint next = ((index + 1) * bits - 1) >> 6;  // java signed >>

            uint shift = bitIndex ^ cell << 6;

            auto [d, lock] = data.borrow();

            // read previous value
            uint64 cellValue = d[cell];
            uint prev = (uint)(d[cell] >> shift & mask);
            
            d[cell] &= ~(mask << shift);                // clear target bits
            d[cell] |= ((uint64)value & mask) << shift; // set target bits

            // Bits overflow into next cell
            if (next != cell)
            {
                uint shift1 = BITS_PER_CELL - shift;
                uint shift2 = bits - shift1;

                uint64 nextValue = d[next];

                prev |= (uint)(d[next] << shift1 & mask);
                // TODO: not 100% on the math here
                d[next] = d[next] >> shift2 << shift2 | ((uint64)value & mask) >> shift1; // last shift is java signed >>
            }

            return prev;
        }

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