#!/usr/bin/env python3
"""P0 enemy concept drafts via the LAN LLM3DGraphics local Turbo image API (DGX3D_AI调用指南.md §4.4).

All five enemies are one 'horn imp' family built as short-legged bipeds (biped_short), so the validated
RockBall enemy rig/motion adapter applies. Drafts go to Doc/Design/Character/Enemy/{id}/drafts/.

Usage:
    python Tools/generate_enemy_designs.py                 # all enemies, 4 drafts each
    python Tools/generate_enemy_designs.py Hornlet Boss    # selected enemies only
"""
from __future__ import annotations

import json
import sys
import time
from pathlib import Path

import requests

BASE = "http://192.168.130.16:8008"
ROOT = Path(__file__).resolve().parents[1]
DESIGN = ROOT / "Doc/Design/Character/Enemy"
STYLE_REFS = [ROOT / "Doc/Design/Character/FireCat/firecat.png", ROOT / "Doc/Design/Character/RockRhino/rockrhino.png"]

COMMON = (
    "Image 1 and image 2 are STYLE ONLY: match their friendly polished chibi 3D mobile-game render, glossy expressive eyes, "
    "rounded appealing volumes and clean readable shapes; do not copy their species, colors or features. "
    "Subject family: small mischievous 'horn imp' monsters that invade a cozy meadow camp in a cute tower-defense game; "
    "cheeky and naughty, never scary, gory or realistic. Base palette violet / lilac skin with cream-ivory horns. "
    "Body plan for rigging: a SHORT-LEGGED BIPED standing upright, two short stubby legs with flat round feet both grounded, "
    "two short separate arms held slightly away from the body, bilaterally symmetrical, neutral calm standing pose, "
    "head facing the viewer. Composition: single full-body character centered, very shallow front three-quarter view "
    "(about 15 degrees off frontal), whole body including horns and feet in frame with generous margin. "
    "Pure solid white #FFFFFF background to all corners, soft studio light, tiny contact shadow under the feet. "
    "Simple smooth matte surfaces suitable for a 2000-triangle mobile model with a 1024 texture. "
    "Exactly one character; no text, no watermark, no UI, no props held in hands, no weapons, no floor base, no scenery, "
    "no particles, no glow, no extra limbs, no wings."
)

ENEMIES = [
    {"id": "Hornlet", "name": "小角怪", "role": "normal", "seed": 61001,
     "subject": "Hornlet, the basic grunt: a small round violet imp, egg-shaped body, one short stubby cream horn on top of the head, "
                "big round eyes and a toothy cheeky grin with one small fang, tiny pointed ears, short tail stub."},
    {"id": "Dashlet", "name": "疾跑怪", "role": "fast", "seed": 62001,
     "subject": "Dashlet, the fast runner: a slimmer, lighter lilac imp with a narrow teardrop body leaning slightly forward, "
                "one sleek swept-back cream horn, long pointed ears swept back like it is in a hurry, eager grinning face, "
                "slightly longer but still short legs with springy feet, thin tail; reads as quick and light."},
    {"id": "Shellhorn", "name": "硬壳怪", "role": "armored", "seed": 63101,
     # v1 (seeds 63001-63004) drifted into four-legged turtles; v2 forces an upright biped with a compact back plate.
     "subject": "Shellhorn, the armored one: an UPRIGHT standing two-legged violet imp, same upright body plan as a small chubby person, "
                "standing on its two short hind legs only, two short arms hanging at its sides, NOT on all fours, not a turtle, not a quadruped. "
                "Stocky wide body; a compact rounded gray-brown stone plate covers its back like a small backpack-sized shell, staying behind the body "
                "and not wider than the shoulders, plus a smooth stone brow plate on its forehead with one short blunt cream horn poking out of it; "
                "stubborn frowning but cute face. The stone plates are solid attached armor with shallow seams, not floating rocks."},
    {"id": "HornCaptain", "name": "角怪队长", "role": "elite", "seed": 64001,
     "subject": "Horn Captain, the elite squad leader: a taller, chunkier deep-violet imp with TWO curved cream horns, a short red scarf "
                "tied around the neck with the knot at the front, confident smug grin with two small fangs, small spiky tuft on the head, "
                "broader shoulders and slightly bigger fists; still short-legged and cute."},
    {"id": "GreatHorn", "name": "大角首领", "role": "boss", "seed": 65001,
     "subject": "Great Horn, the boss: a big heavy dark plum-violet imp chief with a round belly, two large curved ram-like cream horns "
                "and a small crown-like ridge of three tiny horns between them, thick bushy eyebrows, proud bossy grin with fangs, "
                "a short dark purple cape draped over the shoulders, chunky short legs and big round feet; imposing but still cute."},
]


def status(session: requests.Session) -> dict:
    r = session.get(f"{BASE}/api/local-images/status", timeout=20)
    r.raise_for_status()
    return r.json()


def run(enemy: dict, session: requests.Session, design: Path = DESIGN, refs: list[Path] = STYLE_REFS, common: str = COMMON,
        size: tuple[int, int] = (1024, 1024)) -> dict:
    """Generate four drafts for one character into design/{id}/drafts; enemy defaults, overridable for companions."""
    out = design / enemy["id"] / "drafts"
    out.mkdir(parents=True, exist_ok=True)
    prompt = f"{common} {enemy['subject']}"
    data = dict(model="turbo", profile="concept", prompt=prompt, width=size[0], height=size[1], count=4,
                seed=enemy["seed"], transparent="false", negative_prompt="")
    files = [("images", (p.name, p.open("rb"), "image/png")) for p in refs]
    try:
        r = session.post(f"{BASE}/api/local-images/jobs", data=data, files=files, timeout=60)
    finally:
        for _, (_, fh, _) in files:
            fh.close()
    r.raise_for_status()
    job = r.json()
    while job.get("status") not in ("completed", "failed"):
        time.sleep(3)
        job = session.get(f"{BASE}/api/local-images/jobs/{job['id']}", timeout=20).json()
    saved = []
    for img in job.get("images", []):
        path = out / f"{enemy['id']}_{img['seed']}.png"
        d = session.get(f"{BASE}{img['url']}" if img["url"].startswith("/") else img["url"], timeout=60)
        d.raise_for_status()
        path.write_bytes(d.content)
        saved.append(str(path.relative_to(ROOT)).replace("\\", "/"))
    record = {"id": enemy["id"], "name": enemy["name"], "role": enemy["role"], "jobId": job["id"], "status": job["status"],
              "error": job.get("error"), "model": "Qwen/Qwen-Image-2.1-Turbo (LAN LLM3DGraphics local-images)", "profile": "concept",
              "size": list(size), "seedStart": enemy["seed"], "references": [str(p.relative_to(ROOT)).replace("\\", "/") for p in refs],
              "prompt": prompt, "drafts": saved, "generatedUtc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}
    (design / enemy["id"] / "design_drafts.json").write_text(json.dumps(record, ensure_ascii=False, indent=2), encoding="utf-8")
    return record


def main() -> None:
    wanted = set(sys.argv[1:])
    with requests.Session() as session:
        state = status(session)
        if not (state.get("connected") and state.get("ready")):
            raise SystemExit(f"local images not ready: {state}")
        for enemy in ENEMIES:
            if wanted and enemy["id"] not in wanted:
                continue
            rec = run(enemy, session)
            print(json.dumps({"id": rec["id"], "status": rec["status"], "drafts": len(rec["drafts"]), "error": rec["error"]}, ensure_ascii=False))


if __name__ == "__main__":
    main()