"""Regenera imagenes sueltas de un pack ya generado, con semilla nueva.

Recompone el plan con la MISMA semilla (misma pose, camara y stage), regenera
solo los indices pedidos y sustituye el PNG y el registro del manifest. El
fichero viejo se mueve a `_replaced/` (nunca se borra).

    python Engine/regen_pack_items.py --pack <carpeta|nombre> \
        --theme onsen_xilonen --char genshinimpact_xilonen_illustriousXL \
        --items 16,19,21 --seed 42 --new-seed-base 900 --clean-poses

`--seed` es la semilla del PLAN (la que eligio las poses); `--new-seed-base`
es la base de las semillas nuevas de imagen (base + indice).
"""
import argparse
import base64
import json
import random
import shutil
import sys
from pathlib import Path

sys.stdout.reconfigure(encoding="utf-8", errors="backslashreplace")

BASE_DIR = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(BASE_DIR / "Engine"))

import pose_bank  # noqa: E402
import prompt_builder  # noqa: E402
from pack_builder import (FORGE, PACKS_DIR, _image_name, _lora_weight,  # noqa: E402
                          build_run_prompt, compose_theme_jobs, load_affinity,
                          load_cameras, load_pool_lines, parse_stage_quotas,
                          theme_negative)
from stages import STAGES, load_stage_cameras  # noqa: E402

LEVELS = ("sugerente", "spicy", "explicito")


def resolve_pack(arg, char, theme) -> Path:
    p = Path(arg)
    if p.is_dir():
        return p
    for cand in (PACKS_DIR / arg, PACKS_DIR / f"{char}_{theme}"):
        if cand.is_dir():
            return cand
    raise SystemExit(f"[!] no encuentro el pack '{arg}'")


def compose(theme, char, stages, explore, max_variants, clean_poses, seed):
    quotas = parse_stage_quotas(stages)
    target = sum(quotas.values())
    cameras = load_cameras()
    pools = {lv: load_pool_lines(lv) for lv in LEVELS}
    B = prompt_builder.PromptBuilder(str(BASE_DIR))
    char_line = B.get_character(char)["base"]
    rng = random.Random(seed)
    jobs = compose_theme_jobs([theme], target, {}, "per-image", {}, "auto",
                              explore, pose_bank.load_bank(),
                              pose_bank.load_avoid(), pools, cameras, rng,
                              max_variants=max_variants,
                              stage_quotas=quotas,
                              affinity=load_affinity(),
                              stage_cameras=load_stage_cameras(fallback=cameras))
    for job in jobs:
        job.clean_poses = clean_poses
    return jobs, B, char_line


def txt2img(prompt, negative, seed):
    import requests
    payload = {
        "prompt": prompt, "negative_prompt": negative,
        "width": 832, "height": 1216, "steps": 28, "cfg_scale": 5.0,
        "sampler_name": "DPM++ 2M Karras", "batch_size": 1,
        "n_iter": 1, "seed": seed, "clip_skip": 2,
    }
    r = requests.post(f"{FORGE}/sdapi/v1/txt2img", json=payload, timeout=300)
    r.raise_for_status()
    data = r.json()["images"][0]
    if "," in data:
        data = data.split(",", 1)[1]
    return base64.b64decode(data)


def main():
    ap = argparse.ArgumentParser(description="Regenera items sueltos de un pack")
    ap.add_argument("--pack", required=True, help="carpeta del pack o su nombre")
    ap.add_argument("--char", required=True)
    ap.add_argument("--theme", required=True)
    ap.add_argument("--items", required=True, help="indices 1-based: 16,19,21")
    ap.add_argument("--seed", type=int, default=42, help="semilla del PLAN")
    ap.add_argument("--new-seed-base", type=int, default=900,
                    help="base de las semillas nuevas (base + indice)")
    ap.add_argument("--stages", default="5,5,5,5,5")
    ap.add_argument("--explore", type=float, default=0.2)
    ap.add_argument("--max-variants", type=int, default=2)
    ap.add_argument("--clean-poses", action="store_true")
    ap.add_argument("--dry-run", action="store_true")
    a = ap.parse_args()

    run_dir = resolve_pack(a.pack, a.char, a.theme)
    manifest_path = run_dir / "pack_manifest.json"
    records = json.loads(manifest_path.read_text(encoding="utf-8"))
    jobs, B, char_line = compose(a.theme, a.char, a.stages, a.explore,
                                 a.max_variants, a.clean_poses, a.seed)

    wanted = [int(x) for x in a.items.split(",") if x.strip()]
    print(f"[-] pack: {run_dir}")
    print(f"[-] plan: {len(jobs)} jobs (seed {a.seed}) | regenero {wanted}")

    for i in wanted:
        if not 1 <= i <= len(jobs):
            raise SystemExit(f"[!] indice fuera de rango: {i}")
        job = jobs[i - 1]
        prompt = build_run_prompt(B.master_prompt, char_line, job)
        neg = theme_negative(job, B.negative_prompt)
        seed = a.new_seed_base + i
        old = records[i - 1]
        old_name = Path(old["file"]).name
        new_name = _image_name(i, job, seed)
        print(f"\n[{i:02d}] {job.stage} {job.origin} {job.pose_id or ''} "
              f"| {old_name} -> {new_name}")
        print(f"     dress:{job.dress} cam:{job.camera} "
              f"| outfit:{job.outfit_line or '(vacio)'}")
        print("     BREAKs: " + " | ".join(
            s.split(",")[0].strip() + "..." for s in prompt.split("\n\nBREAK\n")))
        if a.dry_run:
            print(f"     prompt: {prompt}")
            continue

        png = txt2img(prompt, neg, seed)
        replaced = run_dir / "raw" / "_replaced"
        replaced.mkdir(parents=True, exist_ok=True)
        for sub in ("raw", "img"):
            src = run_dir / sub / old_name
            if src.exists():
                shutil.move(str(src), str(replaced / old_name))
        for sub in ("raw", "img"):
            (run_dir / sub / new_name).write_bytes(png)

        records[i - 1] = dict(old)
        records[i - 1].update({
            "file": f"img/{new_name}",
            "seed": seed,
            "prompt": prompt,
            "lora_weight": _lora_weight(char_line),
            "regen": {"from": old_name, "from_seed": old.get("seed"),
                      "seed": seed},
        })
        manifest_path.write_text(json.dumps(records, ensure_ascii=False,
                                            indent=1), encoding="utf-8")
        print(f"     [OK] {new_name} (seed {seed})")

    if not a.dry_run:
        print(f"\n[OK] {len(wanted)} imagenes regeneradas en {run_dir}")
        print(f"[i] originales en raw/_replaced/")
        from pack_builder import notify
        notify(f"♻️ Regeneradas {len(wanted)} imágenes de {run_dir.name}")


if __name__ == "__main__":
    main()