"""Crop slots from a labeled screenshot and add them to the real template library.

Usage:
    # 1) preview: crop 10 slots to preview/ so you can see what each slot contains
    python build_library.py samples/shot1.png

    # 2) import: provide 10 comma-separated hero keys (left to right), '?' to skip a slot
    python build_library.py samples/shot1.png tinker,earthshaker,juggernaut,dazzle,vengefulspirit,axe,sniper,slark,lion,drow_ranger

Hero keys must match Steam internal names without the npc_dota_hero_ prefix
(see heroes.json after running fetch_cdn_templates.py).
"""

import json
import sys
import time
from pathlib import Path

import cv2

from common import ROOT, TEMPLATES_REAL, crop_slot, load_config

PREVIEW_DIR = ROOT / "preview"


def known_hero_keys() -> set[str]:
    path = ROOT / "heroes.json"
    if not path.exists():
        return set()
    with open(path, encoding="utf-8") as f:
        return {h["key"] for h in json.load(f)}


def main() -> None:
    if len(sys.argv) < 2:
        sys.exit(__doc__)
    image_path = sys.argv[1]
    labels = sys.argv[2].split(",") if len(sys.argv) > 2 else None

    img = cv2.imread(image_path)
    if img is None:
        sys.exit(f"cannot read image: {image_path}")
    cfg = load_config()
    if not cfg["slots"]:
        sys.exit("config.json has no slots - run calibrate.py first")

    crops = [(slot["index"], crop_slot(img, slot, cfg)) for slot in cfg["slots"]]

    if labels is None:
        PREVIEW_DIR.mkdir(exist_ok=True)
        for idx, crop in crops:
            if crop is not None:
                cv2.imwrite(str(PREVIEW_DIR / f"slot_{idx}.png"), crop)
        print(f"wrote {len(crops)} crops to {PREVIEW_DIR}/ - inspect them, then rerun with labels")
        return

    if len(labels) != 10:
        sys.exit(f"expected 10 labels, got {len(labels)}")
    known = known_hero_keys()
    stamp = time.strftime("%Y%m%d_%H%M%S")
    added = 0
    for (idx, crop), label in zip(crops, labels):
        label = label.strip()
        if label == "?" or crop is None:
            continue
        if known and label not in known:
            print(f"  WARNING slot {idx}: '{label}' not in heroes.json - saved anyway, double-check spelling")
        hero_dir = TEMPLATES_REAL / label
        hero_dir.mkdir(parents=True, exist_ok=True)
        cv2.imwrite(str(hero_dir / f"{stamp}_s{idx}.png"), crop)
        added += 1
    total = sum(1 for _ in TEMPLATES_REAL.rglob("*.png"))
    heroes = sum(1 for d in TEMPLATES_REAL.iterdir() if d.is_dir())
    print(f"added {added} templates; library now {total} images / {heroes} heroes")


if __name__ == "__main__":
    main()