from PIL import Image
from pathlib import Path
from itertools import zip_longest
import re

def read_palette(text):
    palette = dict()
    metadata = dict()
    in_body = False
    for line in text.split('\n'):
        line = line.strip()
        if not line or line == "GIMP Palette":
            continue
        if line == '#':
            in_body = True
            continue
        if in_body:
            *rgb, name = line.split()
            r, g, b = [int(x) for x in rgb]
            index = len(palette)
            palette[f"{r:02x}{g:02x}{b:02x}"] = {"rgb": (r, g, b), "name": name if name.startswith("minecraft:") else f"minecraft:{name}", "index": index}
        else:
            k, *v = line.split(':')
            k, v = k.strip(), ':'.join(v).strip()
            metadata[k] = v
    return {"metadata": metadata, "palette": palette}

def read_imtuple(impath, palette_path):
    palette = read_palette(Path(palette_path).read_text())
    im = Image.open(impath)
    im = im.rotate(180)
    im = im.convert('RGB')
    used_color_hex = tuple(set(f"{r:02x}{g:02x}{b:02x}" for r, g, b in im.getdata()))
    used_mc_palette = tuple(palette["palette"][h]["name"] for h in used_color_hex)
    imdata = (used_color_hex.index(f"{r:02x}{g:02x}{b:02x}") for r, g, b in im.getdata())
    #imtuple = tuple(zip_longest(*([iter(imdata)] * im.width), fillvalue=0))
    #assert len(imtuple) == im.height
    imtuple = imdata
    return (used_mc_palette, imtuple, im.width, im.height)

def read_text(text):
    from mcdraw import text_to_im
    im = text_to_im(text)
    used_mc_palette = ('minecraft:air', 'minecraft:stone')
    imdata = (0 if i else 1 for i in im.getdata())
    imtuple = imdata
    return (used_mc_palette, imtuple, im.width, im.height)

import gzip
from pynbt import NBTFile, TAG_Int, TAG_Compound, TAG_Short, TAG_Byte_Array, TAG_List, TAG_Int_Array
from io import BytesIO
import struct

class VarInt:
    MAX_BYTES = 5
    def __init__(self, io, set_value=None):
        self._io = io
        if set_value:
            self.setval(set_value)
    def setval(self, value):
        ret = bytes()
        while True:
            byte = value & 0x7F
            value >>= 7
            ret += struct.pack("B", byte | (0x80 if value > 0 else 0))
            if value == 0:
                break
        if len(ret) > self.MAX_BYTES:
            raise ValueError("Tried to write too long of a VarInt")
        self._io.write(ret)
    @property
    def value(self):
        number = 0
        bytes_encountered = 0
        while True:
            byte = self._io.read(1)
            if len(byte) < 1:
                raise EOFError("Unexpected end of message.")
            byteord = ord(byte)
            number |= (byteord & 0x7F) << 7 * bytes_encountered
            if not byteord & 0x80:
                break

            bytes_encountered += 1
            if bytes_encountered > self.MAX_BYTES:
                raise ValueError("Tried to read too long of a VarInt")
        return number

def assemble_nbtfile(used_mc_palette, imtuple, width, height):
    blackdataio = BytesIO()
    varint = VarInt(blackdataio)
    for i in imtuple:
        varint.setval(i)
    nbtvalue = {
        'Version': TAG_Int(2),
        'PaletteMax': TAG_Int(len(used_mc_palette)),
        'Palette': TAG_Compound(
            {name: TAG_Int(i) for i, name in enumerate(used_mc_palette)}
        ),
        'Metadata': TAG_Compound({
            'WEOffsetX': TAG_Int(0),
            'WEOffsetY': TAG_Int(0),
            'WEOffsetZ': TAG_Int(0)
        }),
        'Width': TAG_Short(width),
        'Height': TAG_Short(1),
        'Length': TAG_Short(height),
        'DataVersion': TAG_Int(2586),
        'BlockData': TAG_Byte_Array(blackdataio.getvalue()),
        'BlockEntities': TAG_List(TAG_Compound),
        'Offset': TAG_Int_Array([0, 0, 0])
    }
    nbt = NBTFile(value=nbtvalue, name='Schematic')
    return nbt

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='schemgen')
    parser.add_argument('-o', '--output', default='output.schem', help='output worldedit schema file')
    parser.add_argument('-i', '--input', type=str, default="input.png", help='input image')
    parser.add_argument('-t', '--text', type=str, default="", help='input text')
    parser.add_argument('-p', '--palette', type=str, default="minecraft.gpl", help='gimp palette file')
    args = parser.parse_args()

    if args.text:
        nbt = assemble_nbtfile(*read_text(args.text))
    else:
        nbt = assemble_nbtfile(*read_imtuple(args.input, args.palette))
    print(nbt.pretty())
    with gzip.open(args.output, 'w') as f:
        nbt.save(f)