from typing import Optional

import numpy as np
from nbt.nbt import NBTFile, TAG_Compound, TAG_String


class PositionOutOfBounds(ValueError):
    pass


class BlockNotFound(Exception):
    pass


class ChunkBase:
    def __init__(self, x: int, z: int):
        self.x: int = x
        self.z: int = z


class Chunk(ChunkBase):
    def __init__(self, x: int, z: int, nbt: NBTFile):
        super().__init__(x, z)

        self.ox: int = x * 16
        self.oy: int = nbt.get("yPos").value * 16
        self.oz: int = z * 16

        self.nbt: NBTFile = nbt

        self.palettes: list[list[TAG_Compound]] = []
        self.blockstates: list[list[int]] = []
        self.biome_palettes: list[list[TAG_String]] = []
        self.biomes: list[list[int]] = []
        self.get_sections()

    def get_sections(self) -> None:
        self.palettes = []
        self.blockstates = []
        self.biome_palettes = []
        self.biomes = []
        sections = self.nbt.get("sections")
        for s in range(1, len(sections)-1):
            section = sections[s]
            bs_tag = section.get("block_states")
            self.palettes.append(list(bs_tag.get("palette")))
            bs = bs_tag.get("data")
            self.blockstates.append([] if bs is None else list(bs))
            biomes_tag = section.get("biomes")
            self.biome_palettes.append(list(biomes_tag.get("palette")))
            biomes = biomes_tag.get("data")
            self.biomes.append([] if biomes is None else list(biomes))

    def get_block(self, x: int, y: int, z: int) -> Optional[tuple[str, str]]:
        if 0 <= x < 16 and 0 <= z < 16:
            section_i = y // 16 + 4
            oy = y // 16 * 16
            palette = self.palettes[section_i - 1]
            blockstates = self.blockstates[section_i - 1]
            biome_palette = self.biome_palettes[section_i - 1]
            biomes = self.biomes[section_i - 1]
            if blockstates is None or len(blockstates) == 0:
                return palette[0].get("Name").value, biome_palette[0].value

            rel_y = y - oy
            block_bits = max((len(palette) - 1).bit_length(), 4)
            state = self.get_from_long_array(blockstates, block_bits, x, rel_y, z)
            block = palette[state]
            if len(biome_palette) == 1:
                biome_state = 0
            else:
                biome_bits = (len(biome_palette) - 1).bit_length()
                biome_state = self.get_from_long_array(biomes, biome_bits, x // 4, rel_y // 4, z // 4, 4)
            biome = biome_palette[biome_state]
            return block.get("Name").value, biome.value

        else:
            raise PositionOutOfBounds(f"Coordinates x and z should be in range [0:16[")

    def get_from_long_array(self, long_array, bits, x, y, z, size=16) -> int:
        ids_per_long = 64 // bits
        block_i = y * size * size + z * size + x
        long_i = block_i // ids_per_long
        long_val = long_array[long_i]
        if long_val < 0:
            long_val += 2**64

        bit_i = (block_i % ids_per_long) * bits
        value = (long_val >> bit_i) & ((1 << bits) - 1)
        return value

    def get_top_blocks(self) -> tuple[list[list[str]], np.array, np.array, bool]:
        blocks = [[] for _ in range(16)]
        heightmap_shadow = np.zeros((16, 16), dtype="uint16")
        heightmap = np.zeros((16, 16), dtype="uint16")
        biomes = [["minecraft:plains"]*16 for _ in range(16)]
        heightmap_shadow_longs = self.nbt.get("Heightmaps").get("MOTION_BLOCKING")
        heightmap_longs = self.nbt.get("Heightmaps").get("WORLD_SURFACE")
        if heightmap_longs is None:
            return [["minecraft:air"]*16 for _ in range(16)], heightmap_shadow, biomes, True
        heightmap_shadow_longs = heightmap_shadow_longs.value
        heightmap_longs = heightmap_longs.value

        i = 0
        for z in range(16):
            for x in range(16):
                # i = z * 16 + x
                long_i = i // 7
                long1_val = heightmap_shadow_longs[long_i]
                long2_val = heightmap_longs[long_i]
                bit_i = (i % 7) * 9
                height_shadow = (long1_val >> bit_i) & 0b1_1111_1111
                height = (long2_val >> bit_i) & 0b1_1111_1111

                heightmap_shadow[z, x] = height_shadow
                heightmap[z, x] = height
                y = self.oy + height - 1
                if height == 0:
                    block = "minecraft:air"
                    biome = "minecraft:plains"
                else:
                    block, biome = self.get_block(x, y, z)
                blocks[z].append(block)
                biomes[z][x] = biome

                i += 1

        return blocks, heightmap_shadow, biomes, False


class OldChunk(ChunkBase):
    def __init__(self, x: int, z: int, data_version: int):
        super().__init__(x, z)
        self.data_version: int = data_version


class MalformedChunk(ChunkBase):
    pass
