#!/usr/bin/env python3
"""Classify decompressed CNS payloads as images, tile maps, or unknown data."""
from __future__ import annotations

import argparse
import json
from pathlib import Path

from decode_cns import decompress_cns, parse_image


ROOT = Path(__file__).resolve().parents[1]


def classify_payload(path: Path) -> dict:
    item = {
        "name": path.name,
        "decodedSize": None,
        "kind": "unknown",
    }
    try:
        decoded = decompress_cns(path.read_bytes())
    except Exception as exc:
        item["error"] = str(exc)
        return item

    item["decodedSize"] = len(decoded)
    if len(decoded) >= 4:
        width = int.from_bytes(decoded[0:2], "little")
        height = int.from_bytes(decoded[2:4], "little")
        expected_map = 4 + width * height * 4
        if 0 < width <= 256 and 0 < height <= 256 and len(decoded) == expected_map:
            item.update(
                {
                    "kind": "tilemap",
                    "width": width,
                    "height": height,
                    "tileCount": width * height,
                }
            )
            return item
    if len(decoded) >= 2:
        width = decoded[0]
        height = decoded[1]
        expected_map = 2 + width * height * 4
        if 0 < width <= 256 and 0 < height <= 256 and len(decoded) == expected_map:
            item.update(
                {
                    "kind": "tilemap",
                    "width": width,
                    "height": height,
                    "tileCount": width * height,
                    "headerBytes": 2,
                }
            )
            return item

    try:
        width, height, palette, pixels, bpp = parse_image(decoded)
        stride = ((width * bpp + 31) // 32) * 4
        expected_pixels = stride * height
        if 0 < width <= 2048 and 0 < height <= 2048 and len(pixels) >= expected_pixels:
            item.update(
                {
                    "kind": "image",
                    "width": width,
                    "height": height,
                    "bpp": bpp,
                    "paletteColors": len(palette),
                    "pixelBytes": len(pixels),
                }
            )
    except Exception as exc:
        item["error"] = str(exc)

    return item


def markdown(items: list[dict]) -> str:
    counts: dict[str, int] = {}
    for item in items:
        counts[item["kind"]] = counts.get(item["kind"], 0) + 1
    lines = [
        "# CNS Payload Classification",
        "",
        "Generated by `tools/classify_cns_payloads.py` from decompressed CNS payloads.",
        "",
        "## Counts",
        "",
    ]
    for kind, count in sorted(counts.items()):
        lines.append(f"- `{kind}`: {count}")
    lines.extend(
        [
            "",
            "## Tile Maps",
            "",
            "| file | size | decoded bytes |",
            "| --- | ---: | ---: |",
        ]
    )
    for item in items:
        if item["kind"] == "tilemap":
            lines.append(f"| {item['name']} | {item['width']}x{item['height']} | {item['decodedSize']} |")
    lines.extend(
        [
            "",
            "## Unknown",
            "",
            "| file | decoded bytes | note |",
            "| --- | ---: | --- |",
        ]
    )
    for item in items:
        if item["kind"] == "unknown":
            lines.append(f"| {item['name']} | {item.get('decodedSize') or '-'} | {item.get('error', '')} |")
    lines.append("")
    return "\n".join(lines)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--src-dir", type=Path, default=ROOT / "extract_fld")
    parser.add_argument("--json-out", type=Path, default=ROOT / "out" / "cns_payloads.json")
    parser.add_argument("--md-out", type=Path, default=None)
    args = parser.parse_args()

    items = [classify_payload(path) for path in sorted(args.src_dir.glob("*.cns"))]
    args.json_out.parent.mkdir(parents=True, exist_ok=True)
    args.json_out.write_text(json.dumps(items, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    if args.md_out:
        args.md_out.write_text(markdown(items), encoding="utf-8")
    print(f"wrote {len(items)} CNS classifications -> {args.json_out}")
    if args.md_out:
        print(f"wrote markdown -> {args.md_out}")


if __name__ == "__main__":
    main()