from __future__ import annotations

import json
import struct
import sys
from pathlib import Path


ROOT = Path(r"k:/ghidra/Crusader_Decomp")
L0_WDL_PATH = Path(r"e:/emu/psx/Crusader - No Remorse/LSET1/L0.WDL")
GPU_PATH = ROOT / "binary/Crusader - No Remorse (USA) GPU RAM.bin"
OUTPUT_DIR = ROOT / "out/psx_wdl/L0/mode1_live_clut_row_f0_x0"
ROW_BYTES = 2048
LIVE_CLUT_Y = 0xF0
LIVE_CLUT_X = 0

sys.path.insert(0, str(ROOT / "tools"))

from psx_extract_wdl import (
    colorize_indexed_pixels,
    parse_lset_wdl,
    scan_sprite_bundles,
    write_bundle_atlas,
    write_overview_grid,
    write_png_rgba,
)


def main() -> None:
    l0_data = L0_WDL_PATH.read_bytes()
    gpu = GPU_PATH.read_bytes()
    summary = parse_lset_wdl(l0_data)
    if summary is None:
        raise SystemExit("failed to parse L0.WDL")

    region = next(region for region in summary["regions"] if region["name"] == "post_audio_region_04")
    region_data = l0_data[region["offset"] : region["offset"] + region["size"]]
    bundles = scan_sprite_bundles(region_data, max_candidates=160)

    row = gpu[LIVE_CLUT_Y * ROW_BYTES : (LIVE_CLUT_Y + 1) * ROW_BYTES]
    row_words = struct.unpack("<1024H", row)
    palette = list(row_words[LIVE_CLUT_X : LIVE_CLUT_X + 256])

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    entries: list[dict[str, object]] = []
    summary_rows: list[dict[str, object]] = []
    mode1_count = 0

    for bundle in bundles:
        if bundle["mode"] != 1 or not bundle["frames"]:
            continue
        mode1_count += 1
        bundle_dir = OUTPUT_DIR / f"bundle_{bundle['offset']:08X}"
        bundle_dir.mkdir(parents=True, exist_ok=True)
        rendered_frames: list[dict[str, object]] = []
        frame_rows: list[dict[str, object]] = []
        for frame in bundle["frames"]:
            rgba = colorize_indexed_pixels(frame["pixels"], frame["width"], frame["height"], bundle["mode"], palette)
            write_png_rgba(bundle_dir / f"frame_{frame['index']:03d}_live_row_f0_x0.png", rgba, frame["width"], frame["height"])
            rendered_frames.append(
                {
                    "width": frame["width"],
                    "height": frame["height"],
                    "rgba": rgba,
                }
            )
            frame_rows.append(
                {
                    "index": frame["index"],
                    "width": frame["width"],
                    "height": frame["height"],
                    "origin_x": frame["origin_x"],
                    "origin_y": frame["origin_y"],
                    "data_start": frame["data_start"],
                    "consumed": frame["consumed"],
                }
            )
        write_bundle_atlas(bundle_dir / "atlas_live_row_f0_x0.png", rendered_frames)
        metadata = {
            "offset": bundle["offset"],
            "mode": bundle["mode"],
            "palette_formula": "live_gpu_row_0xF0_x0_contiguous_256",
            "palette_source": {
                "gpu_dump": str(GPU_PATH),
                "x": LIVE_CLUT_X,
                "y": LIVE_CLUT_Y,
            },
            "frame_count": bundle["frame_count"],
            "exported_frames": frame_rows,
        }
        (bundle_dir / "palette_formula.json").write_text(json.dumps(metadata, indent=2), encoding="ascii")
        first_frame = bundle["frames"][0]
        first_rgba = rendered_frames[0]["rgba"]
        entries.append(
            {
                "width": first_frame["width"],
                "height": first_frame["height"],
                "rgba": first_rgba,
                "offset": bundle["offset"],
                "area": first_frame["width"] * first_frame["height"],
            }
        )
        summary_rows.append(
            {
                "offset": bundle["offset"],
                "width": first_frame["width"],
                "height": first_frame["height"],
                "frame_count": bundle["frame_count"],
            }
        )

    entries.sort(key=lambda entry: entry["area"], reverse=True)
    overview_entries = [{"width": entry["width"], "height": entry["height"], "rgba": entry["rgba"]} for entry in entries]
    write_overview_grid(OUTPUT_DIR / "overview_live_row_f0_x0.png", overview_entries, columns=4)
    summary_rows.sort(key=lambda row: row["width"] * row["height"], reverse=True)
    (OUTPUT_DIR / "summary.json").write_text(
        json.dumps(
            {
                "palette_formula": "live_gpu_row_0xF0_x0_contiguous_256",
                "mode1_bundle_count": mode1_count,
                "bundles": summary_rows,
            },
            indent=2,
        ),
        encoding="ascii",
    )

    print(f"mode1_bundles={mode1_count}")
    print(f"overview={OUTPUT_DIR / 'overview_live_row_f0_x0.png'}")
    print(f"summary={OUTPUT_DIR / 'summary.json'}")


if __name__ == "__main__":
    main()