SECTIONS = ("PREHEADER", "__lua__", "__gfx__", "__gff__", "__label__", "__map__", "__sfx__", "__music__")

class Pico8Cart(object):
    def __init__(self, sections):
        self._sections = sections

    @classmethod
    def load(cls, fname):
        sections = {}

        with open(fname, "rt") as f:
            section_name = "PREHEADER"
            section_text = []

            def add_section():
                sections[section_name] = "\n".join(section_text)
                section_text.clear()

            for line in f:
                line = line.rstrip("\r\n")
                if line in SECTIONS:
                    add_section()

                    section_name = line
                else:
                    section_text.append(line)

            add_section()
        
        return Pico8Cart(sections)

    def touch(self, section, cb):
        self._sections[section] = touch(section, self._sections.get(section), cb)

    def save(self, fname):
        with open(fname, "wt") as f:
            for s in SECTIONS:
                val = self._sections.get(s)
                val = canonize(s, val)

                if val:
                    if s != "PREHEADER":
                        f.write(f"{s}\n")
                    f.write(f"{val}\n")

def canonize(section, val):
    return touch(section, val, lambda _: ())

def touch(section, val, cb):
    if section in ("__gfx__", "__map__"):
        if val is None: 
            val = ""
        length = 0x2000 if section == "__gfx__" else 0x1000
        row_width=128 if section=="__gfx__" else 256
        msb_first = section == "__map__"
        memory = from_binary(val, length, msb_first)
        cb(memory)
        return to_binary(memory, length, msb_first, row_width)
    return val


def from_binary(pico_data: str, length: int, msb_first: bool):
    hex_data = "".join(pico_data.split("\n"))
    byte_values = []
    for i in range(0,len(hex_data),2):
        x0 = int(hex_data[i],16)
        x1 = int(hex_data[i+1],16)
        byte_values.append(
            (x0 << 4) + x1
            if msb_first else
            (x1 << 4) + x0
        )
    
    assert(len(byte_values) <= length)
    byte_values = (byte_values + [0] * length)[:length]
    return bytearray(byte_values)

def to_binary(memory: bytearray, length: int, msb_first: bool, row_width: int):
    assert(len(memory) == length)

    HEX_CHARS = "0123456789abcdef"

    chars = []
    for i in range(length):
        byte = memory[i]
        msb = (byte & 0xf0) >> 4
        lsb = byte & 0x0f

        if msb_first:
            chars.append(HEX_CHARS[msb])
            chars.append(HEX_CHARS[lsb])
        else:
            chars.append(HEX_CHARS[lsb])
            chars.append(HEX_CHARS[msb])

    lines = []
    for i in range(0, len(chars), row_width):
        lines.append("".join(chars[i:i+row_width]))
    zeroes = "0" * row_width

    while lines and lines[-1] == zeroes:
        lines.pop()
    return "\n".join(lines)