#!/usr/bin/env python3
"""Traces the controller and headset sprites into src/sprite_shapes.hpp.

Usage: python3 scripts/trace_sprites.py SPRITE_DIRECTORY

SPRITE_DIRECTORY holds controller-left.png, controller-right.png, and
headset-front.png, rendered from Valve's CAD files in the logo style: a light
body with transparent gaps and an accent thumbstick. Needs numpy and
opencv-python.
"""
import sys
from pathlib import Path

import cv2
import numpy as np

ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "src" / "sprite_shapes.hpp"
ACCENT = (58, 145, 232)  # kAccent in src/panel.cpp
EPSILON = 0.7  # outline tolerance, in sprite pixels

# (name in the header, file, anchor as a fraction of the image size)
SPRITES = (
    ("kControllerLeft", "controller-left.png", (0.5, 0.29)),
    ("kControllerRight", "controller-right.png", (0.5, 0.29)),
    ("kHeadset", "headset-front.png", (0.5, 0.5)),
)


def rings(mask, height, anchor):
    """Returns the outlines of a mask, scaled so the sprite is 1 unit tall and the anchor is at 0, 0."""
    contours, _ = cv2.findContours(mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
    ax, ay = anchor[0] * mask.shape[1], anchor[1] * mask.shape[0]
    out = []
    for contour in contours:
        if cv2.contourArea(contour) < 4:
            continue
        points = cv2.approxPolyDP(contour, EPSILON, True).reshape(-1, 2)
        if len(points) >= 3:
            out.append([((x + 0.5 - ax) / height, (y + 0.5 - ay) / height) for x, y in points])
    return out


def emit(name, body, accent, width, height):
    lines = [f"inline const Sprite {name}{{{width / height:.4f}f,"]
    for label, group in (("body", body), ("accent", accent)):
        lines.append(f"    // {label}")
        lines.append("    {")
        for ring in group:
            values = ", ".join(f"{x:.4f}f, {y:.4f}f" for x, y in ring)
            lines.append(f"        {{{values}}},")
        lines.append("    },")
    lines.append("};")
    return "\n".join(lines)


def main():
    if len(sys.argv) != 2:
        raise SystemExit(__doc__)
    folder = Path(sys.argv[1])
    parts = []
    for name, file, anchor in SPRITES:
        image = cv2.imread(str(folder / file), cv2.IMREAD_UNCHANGED)
        if image is None or image.shape[2] != 4:
            raise SystemExit(f"{folder / file} is missing or has no alpha channel")
        b, g, r, a = cv2.split(image)
        opaque = a > 127
        color = np.stack([r, g, b], axis=-1).astype(int)
        accent = opaque & (np.abs(color - ACCENT).sum(axis=-1) < 90)
        height, width = a.shape
        body_rings = rings(opaque.astype(np.uint8), height, anchor)
        accent_rings = rings(accent.astype(np.uint8), height, anchor)
        parts.append(emit(name, body_rings, accent_rings, width, height))
        count = sum(len(r) for r in body_rings + accent_rings)
        print(f"{name}: {len(body_rings)} body rings, {len(accent_rings)} accent rings, {count} points")

    OUT.write_text(
        "#pragma once\n\n"
        "// The Steam Frame controllers and headset as filled polygons, for the calibration\n"
        "// animations. Generated by scripts/trace_sprites.py from sprites rendered from\n"
        "// Valve's CAD files. See THIRD_PARTY_NOTICES.md for the license of these outlines.\n\n"
        "#include <vector>\n\n"
        "namespace knockthru::sprites {\n\n"
        "// Points are in sprite heights, so a sprite is 1 unit tall. The anchor sits at 0, 0:\n"
        "// the head center on a controller, and the middle of the headset. Rings use the\n"
        "// even-odd rule. The body fills in the text color and the accent on top of it.\n"
        "struct Sprite {\n"
        "    float aspect; // width over height\n"
        "    std::vector<std::vector<float>> body, accent; // x, y pairs per ring\n"
        "};\n\n"
        + "\n\n".join(parts)
        + "\n\n} // namespace knockthru::sprites\n")
    print(f"Wrote {OUT.relative_to(ROOT)}")


if __name__ == "__main__":
    main()