import sys, os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from game.formats import DATContainer, get_bob_palette, get_world_palette, parse_bob, parse_ico, parse_map
from PIL import Image

PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DAT_PATH = os.environ.get('KELLOGG_DAT', os.path.join(PROJECT_ROOT, 'game', 'original', 'PCKELL.DAT'))
OUT = os.environ.get('KELLOGG_RENDER_OUT', os.path.join(PROJECT_ROOT, 'preview', 'render_out'))
os.makedirs(OUT, exist_ok=True)

container = DATContainer(DAT_PATH)
entries = container.entries

def frame_to_image(frame, palette):
    img = Image.new('RGBA', (frame.width, frame.height), (0, 0, 0, 0))
    px = img.load()
    for y in range(frame.height):
        for x in range(frame.width):
            idx = frame.pixels[y * frame.width + x]
            if idx < 0 or idx > 255:
                continue
            r, g, b = palette[idx]
            px[x, y] = (r, g, b, 255)
    return img

def render_bob(name):
    data = entries[name + '.BOB']
    palette = get_bob_palette(entries, name)
    frames = parse_bob(data)
    d = os.path.join(OUT, 'bob', name)
    os.makedirs(d, exist_ok=True)
    for i, frame in enumerate(frames):
        img = frame_to_image(frame, palette)
        img.save(os.path.join(d, f'{i:03d}.png'))
    # grid sheet (easier to eyeball than a tall vertical stack)
    if frames:
        cols = min(8, len(frames))
        maxw = max(f.width for f in frames)
        maxh = max(f.height for f in frames)
        rows = (len(frames) + cols - 1) // cols
        sheet = Image.new('RGBA', (cols*maxw, rows*maxh), (30,30,30,255))
        for i, frame in enumerate(frames):
            img = frame_to_image(frame, palette)
            x = (i % cols) * maxw
            y = (i // cols) * maxh
            sheet.paste(img, (x, y), img)
        sheet.save(os.path.join(OUT, 'bob', name + '_sheet.png'))
    return len(frames)

if __name__ == '__main__':
    import sys
    if len(sys.argv) > 1:
        name = sys.argv[1]
        n = render_bob(name)
        print(f'{name}.BOB -> {n} frames')
    else:
        print('usage: render_assets.py NAME (basename without extension)')

def render_all_bob():
    results = {}
    for filename in list(entries.keys()):
        if filename.endswith('.BOB'):
            name = filename[:-4]
            try:
                n = render_bob(name)
                results[name] = ('ok', n)
            except Exception as e:
                results[name] = ('ERROR', str(e))
    return results

def tile_to_image(grid, palette):
    img = Image.new('RGBA', (16, 16), (0,0,0,0))
    px = img.load()
    for y in range(16):
        for x in range(16):
            idx = grid[y*16+x]
            if idx < 0 or idx > 255:
                continue
            r,g,b = palette[idx]
            px[x,y] = (r,g,b,255)
    return img

def render_ico(filename):
    # filename e.g. 'W1.ICO', 'COCO_W2.ICO'
    name = filename[:-4]
    if 'W1' in name:
        world = 'W1'
    elif 'W2' in name:
        world = 'W2'
    elif 'W3' in name:
        world = 'W3'
    else:
        raise ValueError(f'cannot determine world for {filename}')
    palette = get_world_palette(entries, world)
    tiles = parse_ico(entries[filename])
    d = os.path.join(OUT, 'ico')
    os.makedirs(d, exist_ok=True)
    cols = 16
    rows = (len(tiles) + cols - 1) // cols
    sheet = Image.new('RGBA', (cols*16, max(rows,1)*16), (40,40,40,255))
    for i, tile in enumerate(tiles):
        img = tile_to_image(tile, palette)
        x = (i % cols) * 16
        y = (i // cols) * 16
        sheet.paste(img, (x,y), img)
    sheet.save(os.path.join(d, name + '.png'))
    return len(tiles)

def render_all_ico():
    results = {}
    for filename in list(entries.keys()):
        if filename.endswith('.ICO'):
            try:
                n = render_ico(filename)
                results[filename] = ('ok', n)
            except Exception as e:
                results[filename] = ('ERROR', str(e))
    return results


_ico_cache = {}
def get_ico_tiles_and_palette(iconame):
    if iconame not in _ico_cache:
        if 'W1' in iconame: world = 'W1'
        elif 'W2' in iconame: world = 'W2'
        elif 'W3' in iconame: world = 'W3'
        else: raise ValueError(iconame)
        palette = get_world_palette(entries, world)
        tiles = parse_ico(entries[iconame])
        _ico_cache[iconame] = (tiles, palette)
    return _ico_cache[iconame]

def render_map(filename):
    # filename e.g. 'W1L0.MAP' -> use W1.ICO + W1.PCC-derived palette
    base = filename[:2]  # 'W1'
    ico_name = base + '.ICO'
    tiles, palette = get_ico_tiles_and_palette(ico_name)
    width, height, cells = parse_map(entries[filename])
    sheet = Image.new('RGBA', (width*16, height*16), (255,0,255,255))
    for i, (tile, ctype) in enumerate(cells):
        x = (i % width) * 16
        y = (i // width) * 16
        if tile < len(tiles):
            img = tile_to_image(tiles[tile], palette)
            sheet.paste(img, (x,y), img)
    d = os.path.join(OUT, 'map')
    os.makedirs(d, exist_ok=True)
    sheet.save(os.path.join(d, filename[:-4] + '.png'))
    return width, height

def render_all_map():
    results = {}
    for filename in list(entries.keys()):
        if filename.endswith('.MAP'):
            try:
                w,h = render_map(filename)
                results[filename] = ('ok', (w,h))
            except Exception as e:
                results[filename] = ('ERROR', str(e))
    return results