"""Small 26.3 world cache used while the legacy automation is being ported.

It retains numeric block-state IDs from chunk palettes and block updates. It
intentionally does not assign 1.16 block meanings to modern state IDs.
"""
from math import floor


class PaletteSection:
    def __init__(self, bits, palette, words, size):
        self.bits = bits
        self.palette = palette
        self.words = words
        self.size = size
        self.overrides = {}
        self.mask = (1 << bits) - 1 if bits else 0
        self.per_word = 64 // bits if bits else 0

    def get(self, index):
        if index in self.overrides:
            return self.overrides[index]
        if self.bits == 0:
            return self.palette[0]
        word = self.words[index // self.per_word]
        palette_index = (word >> ((index % self.per_word) * self.bits)) & self.mask
        return self.palette[palette_index] if self.palette is not None else palette_index

    def set(self, index, block_state_id):
        self.overrides[index] = block_state_id


class ChunkColumn:
    def __init__(self, x, z, sections, min_section):
        self.x = x
        self.z = z
        self.sections = {
            min_section + offset: section
            for offset, section in enumerate(sections)
        }

    def get_block_at(self, x, y, z):
        section = self.sections.get(y // 16)
        if section is None:
            return None
        index = ((y & 15) * 16 + (z & 15)) * 16 + (x & 15)
        return section['blocks'].get(index)

    def set_block_at(self, x, y, z, block_state_id):
        section = self.sections.get(y // 16)
        if section is None:
            return None
        index = ((y & 15) * 16 + (z & 15)) * 16 + (x & 15)
        previous = section['blocks'].get(index)
        section['blocks'].set(index, block_state_id)
        return previous


class WorldState:
    """Sparse loaded-chunk/block-state cache with the old manager's read API."""
    def __init__(self, min_y=-64):
        self.min_y = min_y
        self.min_section = min_y // 16
        self.chunks = {}
        self.index = {}
        from mosfet.info import blocks
        self.indexed_ids = blocks.INDEXED_IDS
        self.center = None
        self.radius = None
        self.dimension = 'minecraft:overworld'
        self.player_id = None

    def set_dimension(self, dimension, dimension_type=None):
        if dimension:
            self.dimension = dimension
        if dimension_type:
            min_y = dimension_type.get('min_y', self.min_y)
            if isinstance(min_y, int):
                self.min_y = min_y
                self.min_section = min_y // 16

    def load_chunk(self, packet):
        self.unload_chunk(packet.x, packet.z)
        chunk = ChunkColumn(packet.x, packet.z, packet.sections,
                            self.min_section)
        self.chunks[(packet.x, packet.z)] = chunk
        for section_y, section in chunk.sections.items():
            palette = section['blocks']
            if not self.indexed_ids.intersection(
                    palette.palette or [palette.get(i) for i in range(palette.size)]):
                continue
            for index in range(palette.size):
                state_id = palette.get(index)
                if state_id not in self.indexed_ids:
                    continue
                local_x = index & 0xF
                local_z = (index >> 4) & 0xF
                local_y = (index >> 8) & 0xF
                pos = (packet.x * 16 + local_x,
                       section_y * 16 + local_y,
                       packet.z * 16 + local_z)
                self.index.setdefault(state_id, set()).add(pos)

    def handle_block_update(self, packet):
        self.set_block_at(packet.x, packet.y, packet.z, packet.block_state_id)

    def handle_section_update(self, packet):
        for x, y, z, state_id in packet.changes:
            self.set_block_at(x, y, z, state_id)

    def get_chunk(self, x, z):
        return self.chunks.get((x, z))

    def get_block_at(self, x, y, z):
        chunk = self.get_chunk(floor(x / 16), floor(z / 16))
        if chunk is None:
            return None
        return chunk.get_block_at(floor(x), floor(y), floor(z))

    def set_block_at(self, x, y, z, block_state_id):
        x, y, z = floor(x), floor(y), floor(z)
        chunk = self.get_chunk(floor(x / 16), floor(z / 16))
        if chunk is None:
            return False
        previous = chunk.set_block_at(x, y, z, block_state_id)
        pos = (x, y, z)
        if previous in self.indexed_ids:
            self.index.get(previous, set()).discard(pos)
        if block_state_id in self.indexed_ids:
            self.index.setdefault(block_state_id, set()).add(pos)
        return True

    def unload_chunk(self, x, z):
        chunk = self.chunks.pop((x, z), None)
        if chunk is None:
            return
        x0, z0 = x * 16, z * 16
        for state_id, locations in list(self.index.items()):
            self.index[state_id] = {
                pos for pos in locations
                if not (x0 <= pos[0] < x0 + 16 and z0 <= pos[2] < z0 + 16)
            }
            if not self.index[state_id]:
                del self.index[state_id]

    def unload_all_chunks(self):
        self.chunks.clear()
        self.index.clear()

    def get_loaded_area(self, ignore_empty=False):
        if not self.chunks:
            return None
        xs = [pos[0] for pos in self.chunks]
        zs = [pos[1] for pos in self.chunks]
        ys = [section_y for chunk in self.chunks.values()
              for section_y in chunk.sections]
        return ((min(xs), min(ys), min(zs)),
                (max(xs), max(ys), max(zs)))

    def check_loaded(self, chunk_distance=None):
        if self.center is None:
            return bool(self.chunks)
        radius = self.radius if chunk_distance is None else chunk_distance
        if radius is None:
            return bool(self.chunks)
        cx, cz = self.center
        return all((x, z) in self.chunks
                   for x in range(cx - radius, cx + radius + 1)
                   for z in range(cz - radius, cz + radius + 1))

    @property
    def loaded_chunk_count(self):
        return len(self.chunks)

    def __repr__(self):
        return ('WorldState(dimension={!r}, chunks={}, min_y={}, center={!r}, '
                'radius={!r})').format(
                    self.dimension, len(self.chunks), self.min_y,
                    self.center, self.radius)


def read_palette_container(stream, *, size, local_bits_max):
    from minecraft.networking.types import Byte, UnsignedLong, VarInt

    bits = Byte.read(stream)
    if bits == 0:
        palette = [VarInt.read(stream)]
    elif bits <= local_bits_max:
        palette_size = VarInt.read(stream)
        palette = [VarInt.read(stream) for _ in range(palette_size)]
    else:
        palette = None

    word_count = ((size + (64 // bits) - 1) // (64 // bits)) if bits else 0
    words = [UnsignedLong.read(stream) for _ in range(word_count)]
    return PaletteSection(bits, palette, words, size)


def decode_chunk_sections(data):
    from io import BytesIO
    from minecraft.networking.types import Short

    stream = BytesIO(data)
    sections = []
    while stream.tell() < len(data):
        start = stream.tell()
        if len(data) - start < 4:
            raise ValueError('truncated chunk section header at byte {}'.format(start))
        non_empty_blocks = Short.read(stream)
        fluid_count = Short.read(stream)
        blocks = read_palette_container(stream, size=4096, local_bits_max=8)
        biomes = read_palette_container(stream, size=64, local_bits_max=3)
        sections.append({
            'non_empty_blocks': non_empty_blocks,
            'fluid_count': fluid_count,
            'blocks': blocks,
            'biomes': biomes,
        })
        if stream.tell() <= start:
            raise ValueError('chunk section decoder made no progress')
    return sections