import zlib
from typing import Optional

import nbt

from src.utils.minecraft.chunk import Chunk, ChunkBase, MalformedChunk, OldChunk


def to_int(bytes_):
    return int.from_bytes(bytes_, byteorder="big")


class Region:
    DATA_VERSION = 3465

    def __init__(self, filepath):
        self.file = open(filepath, "rb")

        self.locations: list[tuple[int, int]] = self.get_locations()
        self.timestamps: list[int] = self.get_timestamps()

        self.chunks = {}

    def get_locations(self) -> list[tuple[int, int]]:
        self.file.seek(0)
        locations = []

        for c in range(1024):
            offset = to_int(self.file.read(3))
            length = to_int(self.file.read(1))
            locations.append((offset, length))

        return locations

    def get_chunk(self, x, z) -> ChunkBase:
        if (x, z) in self.chunks.keys():
            return self.chunks[(x, z)]

        loc = self.locations[(x % 32) + (z % 32) * 32]

        self.file.seek(loc[0] * 4096)
        length = to_int(self.file.read(4))
        compression = to_int(self.file.read(1))
        data = self.file.read(length-1)

        if compression == 2:
            try:
                data = zlib.decompress(data)
            except zlib.error:
                return MalformedChunk(x, z)

        else:
            # print(f"{compression} is not a valid compression type")
            return MalformedChunk(x, z)

        chunk = nbt.nbt.NBTFile(buffer=nbt.chunk.BytesIO(data))
        data_version = chunk.get("DataVersion").value
        if data_version != self.DATA_VERSION:
            print(f"\r    Invalid data version {data_version} for chunk ({x},{z})", end="")
            return OldChunk(x, z, data_version)

        chunk = Chunk(x, z, chunk)
        self.chunks[(x, z)] = chunk

        return chunk

    def get_timestamps(self) -> list[int]:
        return []

    def get_block(self, x: int, y: int, z: int) -> Optional[str]:
        chunk_x, chunk_z = x//16, z//16
        x_rel, z_rel = x % 16, z % 16

        chunk = self.get_chunk(chunk_x, chunk_z)
        if not isinstance(chunk, Chunk):
            return None
        return chunk.get_block(x_rel, y, z_rel)
