#!/usr/bin/env python3
"""Build and verify the approved unit base/action WebP derivatives."""

from __future__ import annotations

import argparse
from concurrent.futures import ProcessPoolExecutor
import json
import os
from pathlib import Path

from PIL import Image, ImageChops, ImageStat, features


ROOT = Path(__file__).resolve().parents[1]
POLICY_PATH = ROOT / "src/game/data/unitActionAssetPolicy.json"
UNIT_DIR = ROOT / "src/assets/images/units"


def load_policy() -> dict:
    return json.loads(POLICY_PATH.read_text(encoding="utf-8"))


def optimized_keys(policy: dict, group: str) -> list[str]:
    if policy.get("optimizeAll") is True:
        if group == "action":
            keys = sorted(path.name.removesuffix("-actions.png") for path in UNIT_DIR.glob("unit-*-actions.png"))
        else:
            keys = sorted(
                path.name.removesuffix(".png")
                for path in UNIT_DIR.glob("unit-*.png")
                if not path.name.endswith("-actions.png")
            )
    else:
        keys = list(policy[f"optimized{group.title()}Keys"])

    expected_count = int(policy.get(f"expected{group.title()}AssetCount", len(keys)))
    if len(keys) != expected_count:
        raise RuntimeError(f"expected {expected_count} {group} sources, found {len(keys)}")
    return keys


def resize_sheet(source: Image.Image, source_size: int, target_size: int, columns: int, rows: int) -> Image.Image:
    expected_size = (columns * source_size, rows * source_size)
    if source.size != expected_size:
        raise ValueError(f"expected {expected_size}, got {source.size}")

    target = Image.new("RGBA", (columns * target_size, rows * target_size))
    for row in range(rows):
        for column in range(columns):
            left = column * source_size
            top = row * source_size
            frame = source.crop((left, top, left + source_size, top + source_size))
            frame = frame.resize((target_size, target_size), Image.Resampling.LANCZOS)
            target.paste(frame, (column * target_size, row * target_size))
    return target


def frame_quality(reference: Image.Image, actual: Image.Image, frame_size: int, columns: int, rows: int) -> dict:
    rgb_error_sum = [0.0, 0.0, 0.0]
    composite_error_sum = [0.0, 0.0, 0.0]
    visible_pixels = 0
    alpha_error_max = 0

    for row in range(rows):
        for column in range(columns):
            box = (
                column * frame_size,
                row * frame_size,
                (column + 1) * frame_size,
                (row + 1) * frame_size,
            )
            expected_frame = reference.crop(box)
            actual_frame = actual.crop(box)
            alpha_diff = ImageChops.difference(expected_frame.getchannel("A"), actual_frame.getchannel("A"))
            alpha_error_max = max(alpha_error_max, alpha_diff.getextrema()[1])

            visible_mask = expected_frame.getchannel("A").point(lambda alpha: 255 if alpha > 8 else 0)
            frame_pixels = ImageStat.Stat(visible_mask).sum[0] / 255
            if frame_pixels <= 0:
                continue

            rgb_mean = ImageStat.Stat(ImageChops.difference(expected_frame, actual_frame), mask=visible_mask).mean[:3]
            for channel in range(3):
                rgb_error_sum[channel] += rgb_mean[channel] * frame_pixels

            for background in ((28, 32, 36, 255), (232, 226, 210, 255)):
                expected_composite = Image.new("RGBA", expected_frame.size, background)
                expected_composite.alpha_composite(expected_frame)
                actual_composite = Image.new("RGBA", actual_frame.size, background)
                actual_composite.alpha_composite(actual_frame)
                composite_mean = ImageStat.Stat(
                    ImageChops.difference(expected_composite.convert("RGB"), actual_composite.convert("RGB"))
                ).mean
                for channel in range(3):
                    composite_error_sum[channel] += composite_mean[channel] * frame_pixels

            visible_pixels += frame_pixels

    divisor = max(1, visible_pixels)
    return {
        "visibleMeanAbsRgb": [round(value / divisor, 3) for value in rgb_error_sum],
        "compositeMeanAbsRgb": [round(value / (divisor * 2), 3) for value in composite_error_sum],
        "alphaErrorMax": alpha_error_max,
    }


def inspect_asset(task: tuple[str, dict, str, bool]) -> dict:
    key, policy, group, write = task
    source_size = int(policy["sourceFrameSize"])
    target_size = int(policy["optimizedFrameSize"])
    columns = int(policy[f"{group}Columns"])
    rows = int(policy[f"{group}Rows"])
    suffix = "-actions" if group == "action" else ""
    source_path = UNIT_DIR / f"{key}{suffix}.png"
    target_path = UNIT_DIR / f"{key}{suffix}.webp"

    source = Image.open(source_path).convert("RGBA")
    resized = resize_sheet(source, source_size, target_size, columns, rows)
    if write:
        target_path.parent.mkdir(parents=True, exist_ok=True)
        resized.save(target_path, "WEBP", lossless=True, method=int(policy["method"]), exact=True)
    if not target_path.exists():
        raise FileNotFoundError(f"missing optimized {group} sheet: {target_path}")

    actual = Image.open(target_path).convert("RGBA")
    if actual.size != resized.size:
        raise ValueError(f"{target_path.name}: expected {resized.size}, got {actual.size}")

    source_bytes = source_path.stat().st_size
    target_bytes = target_path.stat().st_size
    return {
        "group": group,
        "key": key,
        "source": source_path.name,
        "optimized": target_path.name,
        "sourceBytes": source_bytes,
        "optimizedBytes": target_bytes,
        "sourceDecodedBytes": source.width * source.height * 4,
        "optimizedDecodedBytes": actual.width * actual.height * 4,
        "encodedRatio": round(target_bytes / source_bytes, 4),
        **frame_quality(resized, actual, target_size, columns, rows),
    }


def summarize_group(group: str, rows: list[dict]) -> dict:
    group_rows = [row for row in rows if row["group"] == group]
    source_bytes = sum(row["sourceBytes"] for row in group_rows)
    optimized_bytes = sum(row["optimizedBytes"] for row in group_rows)
    source_decoded = sum(row["sourceDecodedBytes"] for row in group_rows)
    optimized_decoded = sum(row["optimizedDecodedBytes"] for row in group_rows)
    if optimized_bytes / source_bytes > 0.55:
        raise RuntimeError(f"optimized {group} sheets did not meet the 45% encoded-size reduction budget")
    if optimized_decoded / source_decoded > 0.38:
        raise RuntimeError(f"optimized {group} sheets did not meet the 62% decoded-memory reduction budget")
    return {
        "assetCount": len(group_rows),
        "sourceMiB": round(source_bytes / 1024 / 1024, 1),
        "optimizedMiB": round(optimized_bytes / 1024 / 1024, 1),
        "encodedReductionPercent": round((1 - optimized_bytes / source_bytes) * 100, 1),
        "sourceDecodedMiB": round(source_decoded / 1024 / 1024, 1),
        "optimizedDecodedMiB": round(optimized_decoded / 1024 / 1024, 1),
        "decodedReductionPercent": round((1 - optimized_decoded / source_decoded) * 100, 1),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--write", action="store_true", help="write the approved WebP derivatives before checking them")
    parser.add_argument("--jobs", type=int, default=min(4, os.cpu_count() or 1), help="parallel conversion/check workers")
    parser.add_argument("--group", choices=("all", "base", "action"), default="all")
    args = parser.parse_args()
    policy = load_policy()

    if not features.check("webp"):
        raise RuntimeError("Pillow was built without WebP support")
    if policy["format"] != "webp" or policy["lossless"] is not True:
        raise ValueError("approved unit derivatives must use lossless WebP")

    groups = ["base", "action"] if args.group == "all" else [args.group]
    tasks = [
        (key, policy, group, args.write)
        for group in groups
        for key in optimized_keys(policy, group)
    ]
    jobs = max(1, min(args.jobs, len(tasks)))
    if jobs == 1:
        rows = [inspect_asset(task) for task in tasks]
    else:
        with ProcessPoolExecutor(max_workers=jobs) as executor:
            rows = list(executor.map(inspect_asset, tasks))

    max_visible_error = max(max(row["visibleMeanAbsRgb"]) for row in rows)
    max_composite_error = max(max(row["compositeMeanAbsRgb"]) for row in rows)
    max_alpha_error = max(row["alphaErrorMax"] for row in rows)
    if max_visible_error != 0 or max_composite_error != 0 or max_alpha_error != 0:
        raise RuntimeError(
            "optimized unit-sheet quality budget failed: "
            f"visible={max_visible_error}, composite={max_composite_error}, alpha={max_alpha_error}"
        )

    report = {
        "policyVersion": policy["version"],
        "sourceFrameSize": policy["sourceFrameSize"],
        "optimizedFrameSize": policy["optimizedFrameSize"],
        "groups": {group: summarize_group(group, rows) for group in groups},
        "quality": {
            "maxVisibleMeanAbsRgb": max_visible_error,
            "maxCompositeMeanAbsRgb": max_composite_error,
            "maxAlphaError": max_alpha_error,
        },
        "largestOptimizedAssets": sorted(rows, key=lambda row: row["optimizedBytes"], reverse=True)[:8],
    }
    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()