from __future__ import annotations

import argparse
import json
from collections import defaultdict
from dataclasses import dataclass, asdict
from pathlib import Path

import numpy as np
from PIL import Image


ROOT = Path(__file__).resolve().parents[1]
UNIT_DIR = ROOT / "src" / "assets" / "images" / "units"
DOCS_DIR = ROOT / "docs"
MANIFEST_PATH = DOCS_DIR / "unit-sprite-rework-manifest.json"
FRAME = 313
BASE_ROWS = 4
BASE_COLS = 16
ACTION_ROWS = 4
ACTION_COLS = 36
BATTLE_TILE_PX = 50

DEFAULT_MANIFEST = {
    "approvedBaseline": [
        "unit-liu-bei",
        "unit-guan-yu",
        "unit-zhang-fei",
        "unit-rebel",
        "unit-rebel-archer",
        "unit-rebel-cavalry",
    ],
    "pendingApproval": ["unit-shu-infantry"],
}


@dataclass
class SheetMetrics:
    key: str
    exists: bool
    size: tuple[int, int] | None = None
    partial_alpha: int | None = None
    opaque_pixels: int | None = None
    border_hits: int | None = None
    min_width: int | None = None
    max_width: int | None = None
    min_height: int | None = None
    max_height: int | None = None
    min_bottom: int | None = None
    max_bottom: int | None = None
    empty_frames: int | None = None
    width_swing: int | None = None
    height_swing: int | None = None
    bottom_swing: int | None = None
    battle_min_width_px: float | None = None
    battle_max_width_px: float | None = None
    battle_min_height_px: float | None = None
    battle_max_height_px: float | None = None
    status: str = "missing"


def measure_sheet(path: Path, key: str, rows: int, cols: int, check_motion_stability = False) -> SheetMetrics:
    if not path.exists():
        return SheetMetrics(key=key, exists=False)

    image = Image.open(path).convert("RGBA")
    alpha = np.array(image.getchannel("A"))
    metrics = SheetMetrics(
        key=key,
        exists=True,
        size=image.size,
        partial_alpha=int(np.count_nonzero((alpha > 0) & (alpha < 255))),
        opaque_pixels=int(np.count_nonzero(alpha == 255)),
        border_hits=0,
        empty_frames=0,
        status="ok",
    )

    expected_size = (FRAME * cols, FRAME * rows)
    widths: list[int] = []
    heights: list[int] = []
    bottoms: list[int] = []
    if image.size != expected_size:
        metrics.status = f"bad-size expected={expected_size}"
        return metrics

    for row in range(rows):
        for col in range(cols):
            frame = alpha[row * FRAME : (row + 1) * FRAME, col * FRAME : (col + 1) * FRAME]
            ys, xs = np.where(frame > 0)
            if len(xs) == 0:
                metrics.empty_frames = (metrics.empty_frames or 0) + 1
                continue

            widths.append(int(xs.max() - xs.min() + 1))
            heights.append(int(ys.max() - ys.min() + 1))
            bottoms.append(int(ys.max()))
            if xs.min() <= 1 or ys.min() <= 1 or xs.max() >= FRAME - 2 or ys.max() >= FRAME - 2:
                metrics.border_hits = (metrics.border_hits or 0) + 1

    if widths:
        metrics.min_width = min(widths)
        metrics.max_width = max(widths)
        metrics.min_height = min(heights)
        metrics.max_height = max(heights)
        metrics.min_bottom = min(bottoms)
        metrics.max_bottom = max(bottoms)
        metrics.width_swing = metrics.max_width - metrics.min_width
        metrics.height_swing = metrics.max_height - metrics.min_height
        metrics.bottom_swing = metrics.max_bottom - metrics.min_bottom
        metrics.battle_min_width_px = round(metrics.min_width * BATTLE_TILE_PX / FRAME, 1)
        metrics.battle_max_width_px = round(metrics.max_width * BATTLE_TILE_PX / FRAME, 1)
        metrics.battle_min_height_px = round(metrics.min_height * BATTLE_TILE_PX / FRAME, 1)
        metrics.battle_max_height_px = round(metrics.max_height * BATTLE_TILE_PX / FRAME, 1)

    issues = []
    if metrics.partial_alpha:
        issues.append(f"partial-alpha={metrics.partial_alpha}")
    if metrics.border_hits:
        issues.append(f"border-hits={metrics.border_hits}")
    if metrics.empty_frames:
        issues.append(f"empty-frames={metrics.empty_frames}")
    if (
        check_motion_stability
        and metrics.max_height is not None
        and metrics.min_height is not None
        and metrics.max_height - metrics.min_height > 48
    ):
        issues.append(f"height-swing={metrics.max_height - metrics.min_height}")
    if (
        check_motion_stability
        and metrics.max_bottom is not None
        and metrics.min_bottom is not None
        and metrics.max_bottom - metrics.min_bottom > 36
    ):
        issues.append(f"bottom-swing={metrics.max_bottom - metrics.min_bottom}")
    metrics.status = "ok" if not issues else "; ".join(issues)
    return metrics


def classify(stem: str) -> str:
    if stem in {"unit-liu-bei", "unit-guan-yu", "unit-zhang-fei"}:
        return "batch1-approved-core"
    if stem in {"unit-rebel", "unit-rebel-archer", "unit-rebel-cavalry"}:
        return "batch1-approved-rebel"
    if stem.startswith("unit-shu-"):
        return "batch2-shu"
    if stem.startswith("unit-wei-"):
        return "batch3-wei"
    if stem.startswith("unit-wu-"):
        return "batch4-wu"
    if stem.startswith("unit-rebel-"):
        return "batch5-rebel-other"
    if stem.startswith("unit-nanman-") or stem == "unit-meng-huo":
        return "batch6-nanman-special"
    return "batch7-unique-officer"


def unit_stems() -> list[str]:
    return sorted(path.stem for path in UNIT_DIR.glob("unit-*.png") if not path.stem.endswith("-actions"))


def load_manifest(path: Path = MANIFEST_PATH) -> dict:
    if not path.exists():
        return DEFAULT_MANIFEST
    manifest = json.loads(path.read_text(encoding="utf-8"))
    return {
        **DEFAULT_MANIFEST,
        **manifest,
    }


def review_stage(stem: str, manifest: dict) -> str:
    if stem in set(manifest.get("approvedBaseline", [])):
        return "approved-baseline"
    if stem in set(manifest.get("pendingApproval", [])):
        return "pending-approval"
    return "queued-rework"


def audit() -> tuple[list[dict], dict[str, int], dict[str, int]]:
    rows: list[dict] = []
    counts: dict[str, int] = defaultdict(int)
    stage_counts: dict[str, int] = defaultdict(int)
    manifest = load_manifest()
    for stem in unit_stems():
        base = measure_sheet(UNIT_DIR / f"{stem}.png", stem, BASE_ROWS, BASE_COLS, check_motion_stability=True)
        actions = measure_sheet(UNIT_DIR / f"{stem}-actions.png", f"{stem}-actions", ACTION_ROWS, ACTION_COLS)
        batch = classify(stem)
        stage = review_stage(stem, manifest)
        counts[batch] += 1
        stage_counts[stage] += 1
        rows.append(
            {
                "batch": batch,
                "review_stage": stage,
                "stem": stem,
                "base": asdict(base),
                "actions": asdict(actions),
            }
        )
    return rows, dict(sorted(counts.items())), dict(sorted(stage_counts.items()))


def write_markdown(rows: list[dict], counts: dict[str, int], stage_counts: dict[str, int], path: Path) -> None:
    lines = [
        "# Unit Sprite Rework Audit",
        "",
        "This file is generated by `scripts/audit-unit-sprite-sheets.py`.",
        "",
        f"Review stages are loaded from `{MANIFEST_PATH.relative_to(ROOT).as_posix()}`.",
        "",
        "Visual review boards are generated by `scripts/render-unit-sprite-review-boards.py`.",
        "",
        "- Review index: `unit-sprite-review-index.png`",
        "- Batch 1 core: `unit-sprite-review-batch1-approved-core.png`",
        "- Batch 1 rebel: `unit-sprite-review-batch1-approved-rebel.png`",
        "- Batch 2 Shu: `unit-sprite-review-batch2-shu.png`",
        "- Batch 3 Wei: `unit-sprite-review-batch3-wei.png`",
        "- Batch 4 Wu: `unit-sprite-review-batch4-wu.png`",
        "- Batch 5 rebel other: `unit-sprite-review-batch5-rebel-other.png`",
        "- Batch 6 Nanman/special: `unit-sprite-review-batch6-nanman-special.png`",
        "- Batch 7 unique officer: `unit-sprite-review-batch7-unique-officer.png`",
        "",
        "## Batch Counts",
        "",
    ]
    for batch, count in counts.items():
        lines.append(f"- {batch}: {count}")

    lines.extend(["", "## Review Stage Counts", ""])
    for stage, count in stage_counts.items():
        lines.append(f"- {stage}: {count}")

    lines.extend(
        [
            "",
            "## Batch Plan",
            "",
            "- batch1-approved-core/rebel: already approved baseline, use as quality reference.",
            "- batch2-shu: Shu generic infantry/spearman/archer/cavalry/strategist/officer variants.",
            "- batch3-wei: Wei generic infantry/archer/cavalry/strategist/officer variants.",
            "- batch4-wu: Wu generic infantry/archer/cavalry/strategist/officer variants.",
            "- batch5-rebel-other: remaining rebel/bandit/yellow-turban variants outside approved batch1.",
            "- batch6-nanman-special: Nanman and special tribal unit variants.",
            "- batch7-unique-officer: unique officer sheets not covered by earlier batches.",
            "",
            "## Review Stage",
            "",
            "- `approved-baseline`: approved hand-painted reference units.",
            "- `pending-approval`: technically prepared, waiting for user approval before expansion/commit/deploy.",
            "- `queued-rework`: listed for later hand-painted rework; technical `ok` does not mean visually approved.",
            "",
            "## Sheet Status",
            "",
            "| batch | stage | stem | base | actions | base bbox | action bbox | battle base px | motion swing |",
            "| --- | --- | --- | --- | --- | --- | --- | --- | --- |",
        ]
    )
    for row in rows:
        base = row["base"]
        actions = row["actions"]
        base_bbox = f'{base.get("min_width")}..{base.get("max_width")} x {base.get("min_height")}..{base.get("max_height")}'
        action_bbox = f'{actions.get("min_width")}..{actions.get("max_width")} x {actions.get("min_height")}..{actions.get("max_height")}'
        battle_base = (
            f'{base.get("battle_min_width_px")}..{base.get("battle_max_width_px")} x '
            f'{base.get("battle_min_height_px")}..{base.get("battle_max_height_px")}'
        )
        motion = (
            f'w{base.get("width_swing")} '
            f'h{base.get("height_swing")} '
            f'b{base.get("bottom_swing")}'
        )
        lines.append(
            f'| {row["batch"]} | {row["review_stage"]} | `{row["stem"]}` | {base["status"]} | {actions["status"]} | {base_bbox} | {action_bbox} | {battle_base} | {motion} |'
        )

    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description="Audit 313px unit sprite sheets.")
    parser.add_argument("--json", type=Path, default=DOCS_DIR / "unit-sprite-rework-audit.json")
    parser.add_argument("--markdown", type=Path, default=DOCS_DIR / "unit-sprite-rework-audit.md")
    args = parser.parse_args()

    rows, counts, stage_counts = audit()
    args.json.write_text(
        json.dumps(
            {
                "manifest": MANIFEST_PATH.relative_to(ROOT).as_posix(),
                "counts": counts,
                "stage_counts": stage_counts,
                "rows": rows,
            },
            indent=2,
        ),
        encoding="utf-8",
    )
    write_markdown(rows, counts, stage_counts, args.markdown)
    print(f"Wrote {args.json}")
    print(f"Wrote {args.markdown}")


if __name__ == "__main__":
    main()