'''
Shared decode library for Kellogg's Tony & Friends asset formats.
Ported (not copied) from the C# reference project
https://github.com/movAX13h/tony-and-friends-in-kelloggs-land (no LICENSE file,
therefore reimplemented from scratch in Python using their format docs as a guide).

Container: PCKELL.DAT holds all 178 assets, footer-indexed (see DATContainer).
'''
import struct
import os
from collections import namedtuple

# ---------------------------------------------------------------------------
# Container (PCKELL.DAT)
# ---------------------------------------------------------------------------

class DATContainer:
    def __init__(self, path):
        with open(path, 'rb') as f:
            self.data = f.read()
        self.entries = {}  # filename -> bytes
        self._parse()

    def _parse(self):
        data = self.data
        num_entries = struct.unpack_from('<H', data, len(data) - 4)[0] + 1
        offset = struct.unpack_from('<i', data, len(data) - 8)[0]
        index_offset = offset

        raw = []
        for _ in range(num_entries):
            name_len = struct.unpack_from('<H', data, offset)[0]
            offset += 2
            filename = data[offset:offset + name_len].decode('ascii', errors='replace')
            offset += name_len
            e_offset = struct.unpack_from('<i', data, offset)[0]
            offset += 4
            raw.append([filename, e_offset])

        for i, e in enumerate(raw):
            length = (raw[i + 1][1] if i + 1 < len(raw) else index_offset) - e[1]
            self.entries[e[0]] = data[e[1]:e[1] + length]

    def extract_all(self, out_dir):
        os.makedirs(out_dir, exist_ok=True)
        for name, payload in self.entries.items():
            with open(os.path.join(out_dir, name), 'wb') as f:
                f.write(payload)


# ---------------------------------------------------------------------------
# Palette (from a .PCC / PCX file's trailing 256-color block)
# ---------------------------------------------------------------------------

def get_pcx_palette(pcc_bytes):
    '''Returns list of 256 (r,g,b) tuples from a PCX v5 file's end-of-file palette.'''
    marker = pcc_bytes[-769]
    if marker != 0x0C:
        raise ValueError('PCX palette marker (0x0C) not found')
    pal_bytes = pcc_bytes[-768:]
    return [tuple(pal_bytes[i:i+3]) for i in range(0, 768, 3)]


def load_pcc(pcc_bytes):
    '''Decode a .PCC (= plain PCX v5, 8bpp, RLE, own trailing 256-color
    palette) via Pillow's battle-tested PCX decoder instead of a hand-rolled
    RLE parser -- see git history 2026-07-22 (commits 5efba4b/62af15e/3410a53)
    for why: TWO from-scratch RLE implementations silently produced visibly
    wrong pixels (shifted/duplicated content) for full-screen 320x200 images
    (KARTE.PCC, KELLOGGS.PCC) despite being "verified" -- Pillow avoids that
    whole class of bug entirely.

    Returns (width, height, indices, palette):
      indices  -- flat list[int] of palette indices, len width*height, row-major
      palette  -- list of 256 (r,g,b) tuples (this file's OWN embedded palette,
                  not any global one)
    '''
    from PIL import Image
    import io
    img = Image.open(io.BytesIO(pcc_bytes))
    img.load()
    if img.mode != 'P':
        raise ValueError(f'expected palette (P) mode PCX, got {img.mode}')
    width, height = img.size
    indices = list(img.getdata())
    raw_pal = img.getpalette()  # flat [r,g,b, r,g,b, ...], may be shorter than 768
    raw_pal = (raw_pal + [0] * 768)[:768]
    palette = [tuple(raw_pal[i:i+3]) for i in range(0, 768, 3)]
    return width, height, indices, palette


def pcc_to_rgba(pcc_bytes, transparent_index=None):
    '''Convenience wrapper around load_pcc(): returns (width, height, rgba_bytes)
    ready for e.g. pygame.image.frombuffer(rgba, (w,h), 'RGBA') or PIL Image.frombytes.
    If transparent_index is given, that palette index becomes alpha=0.'''
    width, height, indices, palette = load_pcc(pcc_bytes)
    out = bytearray(width * height * 4)
    for i, idx in enumerate(indices):
        r, g, b = palette[idx & 0xFF]
        a = 0 if (transparent_index is not None and idx == transparent_index) else 255
        out[i*4:i*4+4] = bytes([r, g, b, a])
    return width, height, bytes(out)


def get_bob_palette(entries, name):
    '''name = basename without extension, e.g. 'TONY' for TONY.BOB.
    Implements the palette-selection rules from Form1.cs (BOB case).'''
    pcc_name = name + '.PCC'
    if pcc_name in entries:
        return list(get_pcx_palette(entries[pcc_name]))
    if len(name) == 1:
        # the ants have names like A.BOB .. O.BOB
        return list(get_pcx_palette(entries['ANTS.PCC']))
    pal = list(get_pcx_palette(entries['W2.PCC']))
    pal[0] = (0, 0, 0)
    return pal


def get_world_palette(entries, world):
    '''world = 'W1'|'W2'|'W3'. Implements ICO/MAP palette rule (addW2Palette):
    W2.PCC carries 16 shared colors (indices 16..31) used for animations/items
    across all worlds, so W1/W3 palettes borrow that slice from W2.'''
    pal = list(get_pcx_palette(entries[world + '.PCC']))
    if world != 'W2':
        w2 = get_pcx_palette(entries['W2.PCC'])
        for i in range(16):
            pal[16 + i] = w2[16 + i]
    return pal


# ---------------------------------------------------------------------------
# BOB (animated sprites: self-modifying x86 'draw code' + embedded pixel data)
# ---------------------------------------------------------------------------

CopyInstr = namedtuple('CopyInstr', ['ega_page', 'offset', 'data'])

def _parse_bob_executable(data):
    '''Interprets the tiny x86 instruction stream used to blit pixels.
    Returns list of CopyInstr, or None if an unrecognized opcode is hit.'''
    instrs = []
    pc = 0
    ega_page = -1
    n = len(data)
    while True:
        if pc >= n:
            return None
        op = data[pc]
        if op == 0x03:  # add si, cx (0x03 0xF1)
            if data[pc+1] != 0xF1:
                return None
            pc += 2
        elif op == 0xCB:  # retf
            pc += 1
            break
        elif op in (0x56, 0x58, 0x5E, 0x50):  # push si / push ax / pop si / push ax
            pc += 1
        elif op == 0xEE:  # out dx, al  -> advance EGA page (cycles 0..3)
            ega_page = (ega_page + 1) & 3
            pc += 1
        elif op == 0xD0:  # rol al, 1  (0xD0 0xC0)
            if data[pc+1] != 0xC0:
                return None
            pc += 2
        elif op == 0x8A:  # mov cl, ah/bl/bh (no-op for our purposes)
            if data[pc+1] not in (0xCC, 0xCB, 0xCF):
                return None
            pc += 2
        elif op == 0xC6:
            sub = data[pc+1]
            if sub == 0x84:  # mov byte ptr [si+AAAA], BB
                offset = data[pc+2] | (data[pc+3] << 8)
                const = data[pc+4]
                pc += 5
                instrs.append(CopyInstr(ega_page, offset, bytes([const])))
            elif sub == 0x44:  # mov byte ptr [si+AA], BB
                offset = data[pc+2]
                const = data[pc+3]
                pc += 4
                instrs.append(CopyInstr(ega_page, offset, bytes([const])))
            else:
                return None
        elif op == 0xC7:
            sub = data[pc+1]
            if sub == 0x44:  # mov word ptr [si+AA], BBBB
                offset = data[pc+2]
                const = data[pc+3] | (data[pc+4] << 8)
                pc += 5
                instrs.append(CopyInstr(ega_page, offset, bytes([const & 0xFF, (const >> 8) & 0xFF])))
            elif sub == 0x84:  # mov word ptr [si+AAAA], BBBB
                offset = data[pc+2] | (data[pc+3] << 8)
                const = data[pc+4] | (data[pc+5] << 8)
                pc += 6
                instrs.append(CopyInstr(ega_page, offset, bytes([const & 0xFF, (const >> 8) & 0xFF])))
            else:
                return None
        else:
            return None
    return instrs


BOB_STRIDE = 84

class BobFrame:
    # pixels: list of palette indices, width*height, -1 = transparent.
    # slot/anchor_x/anchor_y kommen aus dem 14-Byte-Frame-Header (Felder
    # u0/u1/u2, Fund 26.07.2026): slot = Engine-Sprite-Index (0-49 rechts,
    # +50 = Blickrichtung links), anchor = Fusspunkt-Anker relativ zur
    # linken oberen Ecke (anchor_y > height z.B. bei Cocos Hangel-Frames).
    __slots__ = ('width', 'height', 'pixels', 'slot', 'anchor_x', 'anchor_y')
    def __init__(self, width, height, slot=0, anchor_x=0, anchor_y=0):
        self.width = width
        self.height = height
        self.slot = slot
        self.anchor_x = anchor_x
        self.anchor_y = anchor_y
        self.pixels = [-1] * (width * height)

    def set(self, x, y, idx):
        if 0 <= x < self.width and 0 <= y < self.height:
            self.pixels[y * self.width + x] = idx


def parse_bob(data):
    '''Returns list of BobFrame (palette-index pixel grids, -1 = transparent/unset).'''
    frames = []
    pos = 0
    n = len(data)
    while pos < n:
        header = data[pos:pos+14]
        if len(header) < 14:
            raise ValueError(f'truncated header at {pos}')
        slot = struct.unpack_from('<h', header, 0)[0]
        anchor_x = struct.unpack_from('<h', header, 2)[0]
        anchor_y = struct.unpack_from('<h', header, 4)[0]
        width = struct.unpack_from('<h', header, 6)[0]
        height = struct.unpack_from('<h', header, 8)[0]
        ptr_seg_len = struct.unpack_from('<h', header, 10)[0]
        pos += 14

        pointers_data = data[pos:pos+ptr_seg_len]
        pos += ptr_seg_len

        last_instruction = struct.unpack_from('<h', pointers_data, len(pointers_data) - 2)[0]
        next_frame_pos = pos + last_instruction
        exec_segment = data[pos:next_frame_pos]
        pos = next_frame_pos

        instrs = _parse_bob_executable(exec_segment)
        if instrs is None:
            raise ValueError(f'failed to parse exec segment for frame {len(frames)} (len={len(exec_segment)})')

        frame = BobFrame(width, height, slot, anchor_x, anchor_y)
        for instr in instrs:
            for i, byte in enumerate(instr.data):
                p = instr.offset + i
                x = (p % BOB_STRIDE) * 4 + instr.ega_page
                y = p // BOB_STRIDE
                idx = byte - 0x80
                frame.set(x, y, idx)
        frames.append(frame)

    if pos != n:
        raise ValueError(f'missed {n - pos} extra bytes at end of file')

    return frames


# ---------------------------------------------------------------------------
# ICO (16x16 tiles, EGA scrambled pixel order)
# ---------------------------------------------------------------------------

def parse_ico(data):
    '''Returns list of 16x16 palette-index grids (flat lists, len 256).'''
    num_tiles = (len(data) - 1) // 256
    tiles = []
    ptr = 0
    for _ in range(num_tiles):
        grid = [0] * 256
        for y in range(16):
            for page in range(4):
                for x in range(4):
                    k = data[ptr]
                    ptr += 1
                    col = x * 4 + page
                    grid[y * 16 + col] = k - 128
        tiles.append(grid)
    return tiles


# ---------------------------------------------------------------------------
# MAP (grid of tile refs + collision type, big-endian!)
# ---------------------------------------------------------------------------

def parse_map(data):
    if data[0:4] != b'TLE1':
        raise ValueError('MAP signature TLE1 not found')
    width = struct.unpack_from('>h', data, 4)[0]
    height = struct.unpack_from('>h', data, 6)[0]
    unknown = struct.unpack_from('>h', data, 8)[0]
    if unknown != 9:
        raise ValueError(f'unexpected constant {unknown} (expected 9)')

    pos = 10
    cells = []  # row-major, len width*height, each (tile, type)
    for _ in range(width * height):
        value = struct.unpack_from('>H', data, pos)[0]
        pos += 2
        tile = value & 0x1FF
        ctype = value >> 9
        cells.append((tile, ctype))

    if pos != len(data):
        raise ValueError(f'missed {len(data) - pos} extra bytes at end of MAP')

    return width, height, cells