"""
prepare_marine_sprites.py

Convert Marine (Doom player) frames to VB tile format for player 2 rendering.
Source: make_into_4_colors/converted/frames/Marine/ (51 PNGs)

Ideal source: magenta (255,0,255) background for transparency, like the zombie
pipeline (make_into_4_colors_2/convert_to_vb.py). Otherwise background is
detected from corners or most common color; only that background maps to
palette index 0 (transparent), sprite pixels to 1-3.

These frames were extracted sequentially from a Doom player sprite sheet.
Likely layout (5 unique directions per pose, dirs 6-8 mirror 4-2):
  0-4:   Walk A, directions 1-5 (front, front-left, left, back-left, back)
  5-9:   Walk B, directions 1-5
  10-14: Walk C, directions 1-5
  15-19: Walk D, directions 1-5
  20-24: Attack E, directions 1-5
  25-29: Attack F, directions 1-5
  30-34: Pain G, directions 1-5
  35-50: Death H-X (16 frames, direction-independent)

Pipeline:
  1. Convert each frame PNG to VB 4-color palette
  2. Pad to 64x64 (8x8 tiles = 64 chars per frame)
  3. Run grit.exe on each frame
  4. Generate marine_sprites.h with frame pointer table

Output: src/vbdoom/assets/images/sprites/marine/
"""

import os
import subprocess
import shutil
import re
from PIL import Image

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
SOURCE_DIR = os.path.join(SCRIPT_DIR, "make_into_4_colors", "converted", "frames", "Marine")
PROJECT_ROOT = SCRIPT_DIR
GRIT_EXE = os.path.join(PROJECT_ROOT, 'grit_conversions', 'grit.exe')
GRIT_DIR = os.path.join(PROJECT_ROOT, 'grit_conversions')
SPRITES_BASE = os.path.join(PROJECT_ROOT, 'src', 'vbdoom', 'assets', 'images', 'sprites')
OUTPUT_DIR = os.path.join(SPRITES_BASE, 'marine')

GRIT_FLAGS = ['-fh!', '-ftc', '-gB2', '-p!', '-m!']

# VB Palette
VB_PALETTE = [
    (0,   0, 0),    # 0: Black (transparent)
    (85,  0, 0),    # 1: Dark red
    (164, 0, 0),    # 2: Medium red
    (239, 0, 0),    # 3: Bright red
]

TILE_SIZE = 8
TARGET_SIZE = 64  # 8x8 tiles

NUM_FRAMES = 51


def luminance(r, g, b):
    return int(0.299 * r + 0.587 * g + 0.114 * b)


def is_magenta(r, g, b):
    return (r > 220 and g < 30 and b > 220)


def map_to_vb_color(gray):
    if gray < 85:
        return 1
    elif gray < 170:
        return 2
    else:
        return 3


# Sentinel for "magenta background" (same as zombie pipeline: index 0 = transparent).
MAGENTA_BG = (255, 0, 255)


def get_background_color(img, alpha_data):
    """Determine background color: prefer magenta (like zombie Option C), else most common color, else top-left."""
    w, h = img.size
    corners = [(0, 0), (w - 1, 0), (0, h - 1), (w - 1, h - 1)]
    for (x, y) in corners:
        if x < w and y < h:
            a = alpha_data.getpixel((x, y))
            if a < 128:
                return None  # already transparent
            r, g, b = img.getpixel((x, y))
            if is_magenta(r, g, b):
                return MAGENTA_BG
    # Most common color (exact) as fallback; if it dominates, use as background
    color_count = {}
    for y in range(h):
        for x in range(w):
            if alpha_data.getpixel((x, y)) >= 128:
                c = img.getpixel((x, y))
                color_count[c] = color_count.get(c, 0) + 1
    if color_count:
        total = sum(color_count.values())
        most_common = max(color_count.items(), key=lambda t: t[1])
        if most_common[1] > total * 0.20:
            return most_common[0]
    r, g, b = img.getpixel((0, 0))
    return (r, g, b)


def is_background(r, g, b, bg, tolerance=12):
    """True if (r,g,b) is background (transparent). Magenta always counts; else match with small tolerance."""
    if bg is None:
        return False
    if bg == MAGENTA_BG:
        return is_magenta(r, g, b)
    return (abs(r - bg[0]) <= tolerance and
            abs(g - bg[1]) <= tolerance and
            abs(b - bg[2]) <= tolerance)


def enhance_contrast(img):
    """Adaptive contrast enhancement for sprite pixels."""
    pixels = []
    for y in range(img.height):
        for x in range(img.width):
            r, g, b = img.getpixel((x, y))[:3]
            if not is_magenta(r, g, b):
                gray = luminance(r, g, b)
                if gray > 5:
                    pixels.append(gray)
    if not pixels:
        return img
    pixels.sort()
    lo = pixels[int(len(pixels) * 0.02)]
    hi = pixels[int(len(pixels) * 0.98)]
    if hi <= lo:
        return img
    out = img.copy()
    for y in range(img.height):
        for x in range(img.width):
            r, g, b = img.getpixel((x, y))[:3]
            if not is_magenta(r, g, b):
                gray = luminance(r, g, b)
                if gray <= 5:
                    stretched = gray
                else:
                    stretched = int(255.0 * (gray - lo) / (hi - lo))
                    stretched = max(0, min(255, stretched))
                out.putpixel((x, y), (stretched, stretched, stretched))
    return out


def convert_frame_to_vb(input_path):
    """Convert a single frame PNG to 64x64 indexed image (0=transparent, 1-3=reds).
    Returns the indexed PIL Image for tile generation. Background becomes index 0."""
    raw = Image.open(input_path)
    img_rgba = raw.convert('RGBA')
    img = img_rgba.convert('RGB')
    alpha_data = img_rgba.split()[3]
    enhanced = enhance_contrast(img)
    bg = get_background_color(img, alpha_data)

    w, h = img.size

    # Find bounding box of non-background content
    min_x, min_y, max_x, max_y = w, h, 0, 0
    has_content = False
    for y in range(h):
        for x in range(w):
            a = alpha_data.getpixel((x, y))
            if a < 128:
                continue
            r, g, b = img.getpixel((x, y))
            if is_magenta(r, g, b) or (r < 5 and g < 5 and b < 5) or is_background(r, g, b, bg):
                continue
            min_x = min(min_x, x)
            min_y = min(min_y, y)
            max_x = max(max_x, x)
            max_y = max(max_y, y)
            has_content = True

    if not has_content:
        pal_img = Image.new('P', (TARGET_SIZE, TARGET_SIZE), 0)
        pal_img.putpalette([c for i in range(4) for c in VB_PALETTE[i]])
        return pal_img

    cw = max_x - min_x + 1
    ch = max_y - min_y + 1
    scale = min(TARGET_SIZE / cw, TARGET_SIZE / ch)
    if scale > 1.0:
        scale = 1.0
    nw = max(1, int(cw * scale))
    nh = max(1, int(ch * scale))

    cropped = img.crop((min_x, min_y, max_x + 1, max_y + 1))
    enhanced_crop = enhanced.crop((min_x, min_y, max_x + 1, max_y + 1))
    alpha_crop = alpha_data.crop((min_x, min_y, max_x + 1, max_y + 1))
    if scale < 1.0:
        cropped = cropped.resize((nw, nh), Image.LANCZOS)
        enhanced_crop = enhanced_crop.resize((nw, nh), Image.LANCZOS)
        alpha_crop = alpha_crop.resize((nw, nh), Image.LANCZOS)

    out = Image.new('RGB', (TARGET_SIZE, TARGET_SIZE), VB_PALETTE[0])
    off_x = (TARGET_SIZE - nw) // 2
    off_y = TARGET_SIZE - nh  # bottom-aligned

    for y in range(nh):
        for x in range(nw):
            a = alpha_crop.getpixel((x, y))
            if a < 128:
                continue
            r, g, b = cropped.getpixel((x, y))
            if is_magenta(r, g, b) or is_background(r, g, b, bg):
                continue
            er = enhanced_crop.getpixel((x, y))[0]
            idx = map_to_vb_color(er)
            out.putpixel((off_x + x, off_y + y), VB_PALETTE[idx])

    # Build indexed image (0=black/transparent, 1-3=reds) for our own tile generator.
    pal_img = Image.new('P', (TARGET_SIZE, TARGET_SIZE), 0)
    pal_list = [c for i in range(4) for c in VB_PALETTE[i]]
    pal_img.putpalette(pal_list)
    for py in range(TARGET_SIZE):
        for px in range(TARGET_SIZE):
            r, g, b = out.getpixel((px, py))
            idx = 0
            for i in range(4):
                if (r, g, b) == VB_PALETTE[i]:
                    idx = i
                    break
            pal_img.putpixel((px, py), idx)
    return pal_img


def indexed_image_to_vb_tiles(pal_img):
    """Convert 64x64 indexed image (0-3) to VB 2bpp tile data: 64 tiles, 16 bytes each = 256 u32.
    Same format as prepare_particles.py / prepare_keycard_sprites.py: each row = one u16 (LE),
    pixel 0 at bits[1:0] (LSB), then 4 u32 per tile = 2 rows each (row0_lo, row0_hi, row1_lo, row1_hi)."""
    TILE_W, TILE_H = 8, 8
    tiles = []
    for ty in range(8):
        for tx in range(8):
            row_bytes = []
            for row in range(8):
                y = ty * 8 + row
                x0 = tx * 8
                p = [min(3, pal_img.getpixel((x0 + i, y))) for i in range(8)]
                u16_val = sum((p[x] << (x * 2)) for x in range(8))
                row_bytes.append(u16_val & 0xFF)
                row_bytes.append((u16_val >> 8) & 0xFF)
            # 4 u32 per tile: bytes 0-3, 4-7, 8-11, 12-15 (same as bytes_to_u32)
            for i in range(0, 16, 4):
                w = (row_bytes[i] | (row_bytes[i + 1] << 8) |
                     (row_bytes[i + 2] << 16) | (row_bytes[i + 3] << 24))
                tiles.append(w)
    return tiles


def write_marine_c_file(tile_words, output_path, array_name):
    """Write one Marine_NNN.c file with the tile array (256 u32)."""
    lines = [
        "",
        "//{{BLOCK(" + array_name.replace("Tiles", "") + ")",
        "//",
        "// " + array_name + ", 64x64@2, VB 2bpp (index 0 = transparent)",
        "// Generated by prepare_marine_sprites.py",
        "//",
        "//}}BLOCK(" + array_name.replace("Tiles", "") + ")",
        "",
        "const unsigned int " + array_name + "[256] __attribute__((aligned(4)))=",
        "{",
    ]
    for i in range(0, 256, 8):
        row = "\t" + ",".join("0x{:08X}".format(tile_words[i + j]) for j in range(8)) + ","
        lines.append(row)
    lines[-1] = lines[-1].rstrip(",")
    lines.append("};")
    lines.append("")
    with open(output_path, 'w') as f:
        f.write("\n".join(lines))


def run_grit(png_path):
    """Run grit on a single PNG file."""
    cmd = [GRIT_EXE, png_path] + GRIT_FLAGS
    result = subprocess.run(cmd, cwd=GRIT_DIR, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"    ERROR running grit: {result.stderr}")
        return None
    basename = os.path.splitext(os.path.basename(png_path))[0]
    c_file = os.path.join(GRIT_DIR, f'{basename}.c')
    if not os.path.exists(c_file):
        print(f"    ERROR: Expected output not found: {c_file}")
        return None
    return c_file


def get_tiles_array_name(c_file_path):
    """Extract the Tiles array name and size from a grit .c file."""
    with open(c_file_path, 'r') as f:
        content = f.read()
    match = re.search(r'const unsigned int (\w+Tiles)\[(\d+)\]', content)
    if match:
        return match.group(1), int(match.group(2))
    return None, None


def main():
    print("=== Marine Sprite Converter ===\n")

    if not os.path.exists(SOURCE_DIR):
        print(f"ERROR: Source dir not found: {SOURCE_DIR}")
        return

    frame_files = [f"Marine_{i:03d}.png" for i in range(NUM_FRAMES)]
    missing = [f for f in frame_files if not os.path.exists(os.path.join(SOURCE_DIR, f))]
    if missing:
        print(f"WARNING: Missing {len(missing)} frames: {missing[:5]}...")
    frame_files = [f for f in frame_files if os.path.exists(os.path.join(SOURCE_DIR, f))]
    print(f"Found {len(frame_files)} frames to process")

    os.makedirs(OUTPUT_DIR, exist_ok=True)

    TILES_PER_FRAME = 256  # 64 tiles * 4 u32 per tile
    tile_arrays = []

    for i, frame_file in enumerate(frame_files):
        src_path = os.path.join(SOURCE_DIR, frame_file)
        pal_img = convert_frame_to_vb(src_path)
        if pal_img is None:
            continue

        tile_words = indexed_image_to_vb_tiles(pal_img)
        array_name = f"marine_{i:03d}Tiles"
        dest_c = os.path.join(OUTPUT_DIR, f"Marine_{i:03d}.c")
        write_marine_c_file(tile_words, dest_c, array_name)

        tile_arrays.append((array_name, TILES_PER_FRAME, i))

        if (i + 1) % 10 == 0 or i == len(frame_files) - 1:
            print(f"  Converted {i + 1}/{len(frame_files)} frames...")

    # Copy source PNGs for reference
    png_dir = os.path.join(OUTPUT_DIR, 'png')
    os.makedirs(png_dir, exist_ok=True)
    for f in frame_files:
        shutil.copy2(os.path.join(SOURCE_DIR, f), os.path.join(png_dir, f))
    generate_header(tile_arrays, TILES_PER_FRAME)
    print(f"\nDone! {len(tile_arrays)} frames converted.")


def generate_header(tile_arrays, tiles_per_frame):
    """Generate marine_sprites.h with frame pointer table and direction LUTs."""
    header_path = os.path.join(OUTPUT_DIR, "marine_sprites.h")
    guard = "__MARINE_SPRITES_H__"
    bytes_per_frame = tiles_per_frame * 4 if tiles_per_frame else 0

    lines = []
    lines.append(f"#ifndef {guard}")
    lines.append(f"#define {guard}")
    lines.append(f"")
    lines.append(f"/*")
    lines.append(f" * Marine (player 2) sprite frames")
    lines.append(f" * Auto-generated by prepare_marine_sprites.py")
    lines.append(f" *")
    lines.append(f" * Total frames: {len(tile_arrays)}")
    lines.append(f" * Frame size: 64x64 px (8x8 tiles)")
    lines.append(f" * Tiles per frame: {tiles_per_frame} u32 words = {bytes_per_frame} bytes")
    lines.append(f" *")
    lines.append(f" * Frame layout:")
    lines.append(f" *   Walk: front 000,001,002,005 | front-left 006-008,011 | left 012-015 | back-left 018-020,023 | back 024-026,029")
    lines.append(f" *   Shoot: front 003,004 | front-left 009,010 | left 016,017 | back-left 021,022 | back 027,028 (mirror for right)")
    lines.append(f" *   Pain: 030 fl, 031 f, 032 l, 033 b, 034 bl | Death: 035-041 | Gibbed: 042-049")
    lines.append(f" */")
    lines.append(f"")

    for array_name, array_size, idx in tile_arrays:
        lines.append(f"extern const unsigned int {array_name}[{array_size}];")

    lines.append(f"")
    lines.append(f"#define MARINE_FRAME_COUNT {len(tile_arrays)}")
    lines.append(f"#define MARINE_FRAME_BYTES {bytes_per_frame}")
    lines.append(f"#define MARINE_WALK_ANIM_FRAMES  4")
    lines.append(f"#define MARINE_ATTACK_ANIM_FRAMES 2")
    lines.append(f"#define MARINE_DEATH_FRAMES 7")
    lines.append(f"#define MARINE_GIB_FRAMES 8")
    lines.append(f"")

    lines.append(f"static const unsigned int* const MARINE_FRAMES[{len(tile_arrays)}] = {{")
    for i, (array_name, array_size, idx) in enumerate(tile_arrays):
        comma = "," if i < len(tile_arrays) - 1 else ""
        lines.append(f"    {array_name}{comma}")
    lines.append(f"}};")
    lines.append(f"")

    # Walk/attack/pain/death/gib LUTs (frame indices from source Marine_000.png .. Marine_050.png)
    lines.append(f"/* Walk frames: [direction][pose] -> frame index */")
    lines.append(f"static const u8 MARINE_WALK_FRAMES[8][4] = {{")
    lines.append(f"    {{  0,  1,  2,  5}},   /* dir 0: front */")
    lines.append(f"    {{  6,  7,  8, 11}},   /* dir 1: front-left */")
    lines.append(f"    {{ 12, 13, 14, 15}},   /* dir 2: left */")
    lines.append(f"    {{ 18, 19, 20, 23}},   /* dir 3: back-left */")
    lines.append(f"    {{ 24, 25, 26, 29}},   /* dir 4: back */")
    lines.append(f"    {{ 18, 19, 20, 23}},   /* dir 5: back-right (mirror 3) */")
    lines.append(f"    {{ 12, 13, 14, 15}},   /* dir 6: right (mirror 2) */")
    lines.append(f"    {{  6,  7,  8, 11}},   /* dir 7: front-right (mirror 1) */")
    lines.append(f"}};")
    lines.append(f"")

    lines.append(f"/* Attack frames: [direction][pose] -> frame index */")
    lines.append(f"static const u8 MARINE_ATTACK_FRAMES[8][2] = {{")
    lines.append(f"    {{ 3,  4}},   /* dir 0: front */")
    lines.append(f"    {{ 9, 10}},   /* dir 1: front-left */")
    lines.append(f"    {{16, 17}},   /* dir 2: left */")
    lines.append(f"    {{21, 22}},   /* dir 3: back-left */")
    lines.append(f"    {{27, 28}},   /* dir 4: back */")
    lines.append(f"    {{21, 22}},   /* dir 5: back-right (mirror 3) */")
    lines.append(f"    {{16, 17}},   /* dir 6: right (mirror 2) */")
    lines.append(f"    {{ 9, 10}},   /* dir 7: front-right (mirror 1) */")
    lines.append(f"}};")
    lines.append(f"")

    lines.append(f"/* Pain frames: [direction] -> frame index (030 fl, 031 f, 032 l, 033 b, 034 bl) */")
    lines.append(f"static const u8 MARINE_PAIN_FRAMES[8] = {{31, 30, 32, 34, 33, 34, 32, 30}};")
    lines.append(f"")

    lines.append(f"/* Death frames: 035-041 (7 frames) */")
    lines.append(f"static const u8 MARINE_DEATH_FRAMES_LUT[7] = {{35, 36, 37, 38, 39, 40, 41}};")
    lines.append(f"")

    lines.append(f"/* Gibbed frames: 042-049 (8 frames) */")
    lines.append(f"static const u8 MARINE_GIB_FRAMES_LUT[8] = {{42, 43, 44, 45, 46, 47, 48, 49}};")
    lines.append(f"")

    lines.append(f"#endif /* {guard} */")
    lines.append(f"")

    with open(header_path, 'w') as f:
        f.write('\n'.join(lines))
    print(f"  Header generated: marine_sprites.h")


if __name__ == '__main__':
    main()
