magicciv/tools/sprite-generation/engine/starter.py

"""Register and track the curated Game-1 starter sprite set."""
from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path

from engine.prompts import (
    compose_prompt,
    get_generation_size,
    get_negative,
    get_target_size,
)
from engine.registry import SpriteRegistry

TOOL_DIR = Path(__file__).resolve().parents[1]
PROJECT = TOOL_DIR.parent.parent
RESOURCES = PROJECT / "public" / "resources"
MANIFEST_PATH = TOOL_DIR / "starter_manifest.json"
WILDS_PATH = RESOURCES / "wilds" / "wilds.json"
WONDERS_PATH = RESOURCES / "tiles" / "water_and_wonders.json"
LAND_SPECIAL_PATH = RESOURCES / "tiles" / "land_special.json"


@dataclass
class StarterReport:
    registered: list[str] = field(default_factory=list)
    reset: list[str] = field(default_factory=list)
    missing_data: list[str] = field(default_factory=list)
    skipped: list[str] = field(default_factory=list)


def load_manifest(path: Path | None = None) -> dict:
    return json.loads((path or MANIFEST_PATH).read_text(encoding="utf-8"))


def starter_sprite_ids(manifest: dict | None = None) -> list[str]:
    """Flat list of registry sprite ids for the starter set."""
    m = manifest or load_manifest()
    ids: list[str] = []
    for uid in m.get("units", []):
        ids.append(f"units/{uid}")
    for bid in m.get("buildings", []):
        ids.append(f"buildings/{bid}")
    for lid in m.get("landmarks", []):
        ids.append(f"landmarks/{lid}")
    for lid in m.get("lairs", []):
        ids.append(f"lairs/{lid}")
    return ids


def _load_building(entity_id: str) -> dict | None:
    path = RESOURCES / "buildings" / f"{entity_id}.json"
    if not path.exists():
        return None
    raw = json.loads(path.read_text(encoding="utf-8"))
    if isinstance(raw, list):
        for item in raw:
            if isinstance(item, dict) and item.get("id") == entity_id:
                return {**item, "entity_id": entity_id}
        return None
    if isinstance(raw, dict):
        return {**raw, "entity_id": entity_id}
    return None


def _load_landmark(entity_id: str) -> dict | None:
    for source in (WONDERS_PATH, LAND_SPECIAL_PATH):
        if not source.exists():
            continue
        raw = json.loads(source.read_text(encoding="utf-8"))
        terrains = raw.get("terrains", raw.get("tiles", []))
        if isinstance(raw, dict):
            for value in raw.values():
                if isinstance(value, list):
                    terrains = value
                    break
        for entry in terrains:
            if isinstance(entry, dict) and entry.get("id") == entity_id:
                return {**entry, "entity_id": entity_id}
    return None


def _load_lair(entity_id: str) -> dict | None:
    if not WILDS_PATH.exists():
        return None
    raw = json.loads(WILDS_PATH.read_text(encoding="utf-8"))
    lair_types = raw.get("wilds", raw).get("lair_types", [])
    for entry in lair_types:
        if entry.get("id") == entity_id:
            return {**entry, "entity_id": entity_id}
    return None


def _unit_base_id(entity_id: str) -> str:
    base = entity_id
    for suffix in ("_dwarves_m", "_dwarves_f", "_m", "_f"):
        if base.endswith(suffix):
            base = base[: -len(suffix)]
            break
    return base


def _load_unit_entity(entity_id: str) -> dict:
    """Best-effort unit entity_data for prompt composition."""
    base_id = _unit_base_id(entity_id)
    for base in (PROJECT / "public" / "games" / "age-of-dwarves" / "data" / "units",
                 RESOURCES / "units"):
        if not base.exists():
            continue
        direct = base / f"{base_id}.json"
        if direct.exists():
            data = json.loads(direct.read_text(encoding="utf-8"))
            if isinstance(data, dict):
                return {**data, "entity_id": entity_id}
        for path in base.glob("*.json"):
            raw = json.loads(path.read_text(encoding="utf-8"))
            if isinstance(raw, list):
                for item in raw:
                    if isinstance(item, dict) and item.get("id") == base_id:
                        return {**item, "entity_id": entity_id}
            elif isinstance(raw, dict) and raw.get("id") == base_id:
                return {**raw, "entity_id": entity_id}
    combat = "civilian" if base_id in ("worker", "founder") else "melee"
    return {
        "entity_id": entity_id,
        "name": base_id.replace("_", " ").title(),
        "description": "",
        "combat_type": combat,
        "keywords": [],
    }


def register_starter_set(
    registry: SpriteRegistry,
    *,
    manifest: dict | None = None,
    reset_status: bool = True,
) -> StarterReport:
    """Ensure every starter sprite exists in spritegen.db and is queued."""
    m = manifest or load_manifest()
    report = StarterReport()

    def _upsert(
        sprite_id: str,
        category: str,
        entity_id: str,
        entity_data: dict,
        install_path: str,
        source_file: str,
    ) -> None:
        prompt = compose_prompt(category, entity_data)
        negative = get_negative(category, combat_type=entity_data.get("combat_type", ""))
        gen_w, gen_h = get_generation_size(category)
        tgt_w, tgt_h = get_target_size(category)

        existed = registry.conn.execute(
            "SELECT status FROM sprites WHERE id=?", (sprite_id,),
        ).fetchone()

        registry.upsert_sprite(
            id=sprite_id,
            category=category,
            entity_id=entity_id,
            prompt=prompt,
            negative_prompt=negative,
            gen_width=gen_w,
            gen_height=gen_h,
            target_width=tgt_w,
            target_height=tgt_h,
            source_file=source_file,
            install_path=install_path,
        )
        if reset_status or not existed:
            registry.update_sprite_status(sprite_id, "needed")
            report.reset.append(sprite_id)
        report.registered.append(sprite_id)

    for uid in m.get("units", []):
        sprite_id = f"units/{uid}"
        entity_data = _load_unit_entity(uid)
        _upsert(
            sprite_id, "units", uid, entity_data,
            f"sprites/units/{uid}.png", "starter_manifest.json",
        )

    for bid in m.get("buildings", []):
        entity_data = _load_building(bid)
        if not entity_data:
            report.missing_data.append(f"buildings/{bid}")
            continue
        _upsert(
            f"buildings/{bid}", "buildings", bid, entity_data,
            f"sprites/buildings/{bid}.png", f"resources/buildings/{bid}.json",
        )

    for lid in m.get("landmarks", []):
        entity_data = _load_landmark(lid)
        if not entity_data:
            report.missing_data.append(f"landmarks/{lid}")
            continue
        _upsert(
            f"landmarks/{lid}", "landmarks", lid, entity_data,
            f"sprites/terrain/{lid}.png", "resources/tiles/",
        )

    for lid in m.get("lairs", []):
        entity_data = _load_lair(lid)
        if not entity_data:
            report.missing_data.append(f"lairs/{lid}")
            continue
        _upsert(
            f"lairs/{lid}", "lairs", lid, entity_data,
            f"sprites/lairs/{lid}.png", "resources/wilds/wilds.json",
        )

    return report