"""
batch_poses_chars: batch de personajes distintos x poses distintas x entornos
distintos, con la estructura del piloto "Diferent poses".

Estructura (identica al piloto de referencia):
  - 3 imagenes por personaje: 1 sugerente + 1 spicy + 1 explicito
  - camaras en el mismo orden: from above / from front / from side
  - un tema (entorno) distinto por personaje
  - salida consolidada por personaje: <run>/<clave>/img/*.png + manifest

Regla dura: NINGUNA pose se repite en todo el lote. Las poses se reparten a mano
sin reemplazo desde el banco y se fijan con `bank_ids`; no se deja al azar de
`compose_theme_jobs` porque su contador de reuso es por personaje, no global.

La semilla es aleatoria por imagen (`-1`): no se fija ninguna.

    python Engine/batch_poses_chars.py --chars 12 --dry-run
    python Engine/batch_poses_chars.py --chars 12
    python Engine/batch_poses_chars.py --chars 12 --exclude-used --pick-seed 20261028
    python Engine/batch_poses_chars.py --keys Makima_v2,Furina --themes oficina,playa

Salida: Output/pilots/Diferent poses/BATCH_<stamp>/<clave>/img/*.png
"""
import argparse
import json
import random
import re
import sys
from datetime import datetime
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"))
sys.path.insert(0, str(BASE_DIR / "Archivos"))

import pose_bank  # noqa: E402
import pack_builder as pb  # noqa: E402
from prompt_builder import PromptBuilder  # noqa: E402

OUT_ROOT = BASE_DIR / "Output" / "pilots" / "Diferent poses"
LEVELS = ("sugerente", "spicy", "explicito")
QUOTAS = {"sugerente": 1, "spicy": 1, "explicito": 1}
# tema ligado a un personaje concreto: no vale para un lote multitudinario
THEMES_BLOCKED = {"onsen_xilonen"}


def usable_poses(bank, avoid, skip_ids=None) -> dict:
    """{nivel: [entradas]} sin las blacklisteadas (avoid.jsonl) ni las vetadas."""
    skip = {x.strip() for x in (skip_ids or []) if x.strip()}
    out = {}
    for lv in LEVELS:
        out[lv] = [e for e in bank
                   if e.get("nivel") == lv
                   and e.get("id") not in skip
                   and not pose_bank.is_avoided(e.get("pose_line", ""), avoid)]
    return out


def used_characters(root=None) -> set:
    """Claves ya usadas en cualquier lote/pack de `Output/`.

    Mira dos cosas, porque no todos los runs dejan el mismo artefacto:
      - `plan.json` (batches: `plan[].char`)
      - `pack_manifest.json` (packs: `char` en cada record)
    """
    root = Path(root) if root else BASE_DIR / "Output"
    used = set()
    for p in root.rglob("plan.json"):
        try:
            doc = json.loads(p.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            continue
        for item in (doc.get("plan") or []):
            if isinstance(item, dict) and item.get("char"):
                used.add(item["char"])
    for p in root.rglob("pack_manifest.json"):
        try:
            recs = json.loads(p.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            continue
        if isinstance(recs, dict):
            recs = [recs]
        for rec in recs:
            if isinstance(rec, dict) and rec.get("char"):
                used.add(rec["char"])
    return used


def pick_characters(B, count, rng, keys=None, exclude=None) -> list:
    """Personajes con LoRA en disco (o los pedidos a mano). -> [(clave, linea)]"""
    if keys:
        out = []
        for k in keys:
            c = B.get_character(k)
            if not c:
                raise SystemExit(f"[!] personaje '{k}' no esta en personajes.txt")
            out.append((k, c["base"]))
        return out

    from character_keys import lora_stems
    available = lora_stems()
    exclude = exclude or set()
    cands = []
    dropped = 0
    skipped = 0
    for k in B.get_character_names():
        c = B.get_character(k)
        if not c:
            continue
        if k in exclude:
            skipped += 1
            continue
        m = re.search(r"<lora:([^:>]+):", c["base"] or "")
        if not m or not available or m.group(1).lower() not in available:
            dropped += 1
            continue
        cands.append((k, c["base"]))
    print(f"[-] personajes con LoRA en disco: {len(cands)} "
          f"({dropped} descartados por LoRA ausente o sin <lora:>, "
          f"{skipped} ya usados)")
    rng.shuffle(cands)
    if count > len(cands):
        print(f"[!] pedias {count} personajes y solo hay {len(cands)} con LoRA")
    return cands[:count]


def pick_themes(count, rng, wanted=None) -> list:
    """Un tema distinto por personaje, sin temas ligados a un personaje."""
    themes = pb.load_themes()
    if wanted:
        for t in wanted:
            if t not in themes:
                raise SystemExit(f"[!] tema '{t}' no existe en {pb.THEMES_FILE}")
        pool = list(wanted)
    else:
        pool = [t for t in themes if t not in THEMES_BLOCKED]
        rng.shuffle(pool)
    if count > len(pool):
        print(f"[!] pedias {count} temas y solo hay {len(pool)} disponibles")
    # ciclo por si pidieras mas personajes que temas (avisa y repite)
    return [pool[i % len(pool)] for i in range(count)]


def deal(chars, themes, poses_by_level, rng) -> list:
    """Reparte poses sin reemplazo: 1 por nivel y personaje.

    Devuelve [{char, theme, poses: {nivel: entrada}}]. El cuello de botella es
    `sugerente` (16 usables), asi que el tamano maximo del lote lo marca ese
    nivel.
    """
    limit = min(len(poses_by_level[lv]) for lv in LEVELS)
    if len(chars) > limit:
        print(f"[!] el banco solo permite {limit} personajes sin repetir pose "
              f"(sugerente={len(poses_by_level['sugerente'])}, "
              f"spicy={len(poses_by_level['spicy'])}, "
              f"explicito={len(poses_by_level['explicito'])}); recorto a {limit}")
        chars = chars[:limit]
        themes = themes[:limit]

    plan = []
    for i, (key, line) in enumerate(chars):
        plan.append({
            "char": key,
            "char_line": line,
            "theme": themes[i],
            "poses": {lv: poses_by_level[lv][i] for lv in LEVELS},
        })
    return plan


def check_no_repeat(plan) -> None:
    ids = [item["poses"][lv]["id"] for item in plan for lv in LEVELS]
    dupes = sorted({x for x in ids if ids.count(x) > 1})
    if dupes:
        raise SystemExit(f"[!] poses repetidas en el plan: {dupes}")
    print(f"[-] poses unicas: {len(ids)}/{len(ids)} (sin repeticion)")


def load_done(char_dir, ids) -> list | None:
    """Records de un personaje ya generado, si su manifest cuadra con el plan.

    Sirve a `--resume`: lo creado se queda creado, no se regenera. Devuelve None
    si falta el manifest, si las poses no coinciden o si falta algun PNG.
    """
    man = Path(char_dir) / "pack_manifest.json"
    if not man.exists():
        return None
    try:
        records = json.loads(man.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None
    if not isinstance(records, list) or len(records) != len(ids):
        return None
    if [r.get("pose_id") for r in records] != list(ids):
        return None
    for rec in records:
        if not (Path(char_dir) / rec["file"]).exists():
            return None
    return records


def write_contact(paths, out_path, cols=3) -> None:
    """Hoja de contactos del lote (si hay PIL)."""
    try:
        from PIL import Image
    except ImportError:
        return
    thumbs = []
    for p in paths:
        try:
            im = Image.open(p)
            im.thumbnail((420, 620))
            thumbs.append(im.convert("RGB"))
        except Exception:
            continue
    if not thumbs:
        return
    w, h = thumbs[0].size
    rows = (len(thumbs) + cols - 1) // cols
    sheet = Image.new("RGB", (cols * w, rows * h), "white")
    for i, im in enumerate(thumbs):
        sheet.paste(im, ((i % cols) * w, (i // cols) * h))
    sheet.save(out_path, quality=88)


def main():
    p = argparse.ArgumentParser(
        description="Batch de personajes x poses x entornos (estructura del "
                    "piloto 'Diferent poses')")
    p.add_argument("--chars", type=int, default=12,
                   help="cuantos personajes (1 por tema y 3 imagenes); "
                        "maximo real = poses usables de sugerente (16)")
    p.add_argument("--keys", default=None,
                   help="personajes concretos separados por coma (claves de personajes.txt)")
    p.add_argument("--exclude-used", action="store_true",
                   help="no repite personajes ya usados en cualquier lote/pack de Output/")
    p.add_argument("--used-root", default=None,
                   help="carpeta donde buscar los lotes previos (default Output/)")
    p.add_argument("--themes", default=None,
                   help="temas concretos separados por coma (default: al azar, distintos)")
    p.add_argument("--exclude-poses", default=None,
                   help="ids de pose a vetar en este lote, separados por coma "
                        "(p.ej. pb0009)")
    p.add_argument("--seed", type=int, default=None,
                   help="semilla FIJA para todas las imagenes (default: aleatoria "
                        "por imagen, como -1); fijala solo para reproducir un lote")
    p.add_argument("--pick-seed", type=int, default=None,
                   help="semilla del reparto (personajes/temas/poses), para que el "
                        "plan sea reproducible; no afecta a la generacion")
    p.add_argument("--out", default=None, help="carpeta de salida")
    p.add_argument("--budget", type=int, default=None,
                   help="maximo de intentos de generacion (default 3 por imagen)")
    p.add_argument("--dry-run", action="store_true",
                   help="imprime el plan y los prompts sin llamar a Forge")
    p.add_argument("--resume", action="store_true",
                   help="no regenera los personajes que ya tienen manifest "
                        "completo en la carpeta de salida (lo creado, creado esta)")
    args = p.parse_args()

    rng = random.Random(args.pick_seed)

    bank = pose_bank.load_bank()
    avoid = pose_bank.load_avoid()
    skip = ([x.strip() for x in args.exclude_poses.split(",")]
            if args.exclude_poses else None)
    poses_by_level = usable_poses(bank, avoid, skip)
    for lv in LEVELS:
        print(f"[-] poses {lv}: {len(poses_by_level[lv])} usables")
        rng.shuffle(poses_by_level[lv])

    B = PromptBuilder(str(BASE_DIR))
    keys = ([x.strip() for x in args.keys.split(",") if x.strip()]
            if args.keys else None)
    used = used_characters(args.used_root) if args.exclude_used else None
    if keys and used:
        clash = sorted(set(keys) & used)
        if clash:
            raise SystemExit(f"[!] --exclude-used y ya estaban usados: {clash}")
    chars = pick_characters(B, args.chars, rng, keys, used)
    if not chars:
        raise SystemExit("[!] no hay personajes disponibles")
    themes = pick_themes(len(chars), rng,
                         [x.strip() for x in args.themes.split(",") if x.strip()]
                         if args.themes else None)

    plan = deal(chars, themes, poses_by_level, rng)
    check_no_repeat(plan)

    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    if args.out:
        run_dir = Path(args.out)
    elif args.resume:
        # sin --out: reanudar el ultimo BATCH_ de la carpeta de pilotos
        done = sorted(OUT_ROOT.glob("BATCH_*"), key=lambda p: p.name)
        if not done:
            raise SystemExit(f"[!] --resume sin --out y no hay BATCH_* en {OUT_ROOT}")
        run_dir = done[-1]
    else:
        run_dir = OUT_ROOT / f"BATCH_{stamp}"

    print(f"[-] {len(plan)} personajes x 3 niveles = {len(plan) * 3} imagenes")
    print(f"[-] salida: {run_dir}")
    for item in plan:
        print(f"    {item['char']:<28} tema {item['theme']:<12} "
              + " | ".join(f"{lv}:{item['poses'][lv]['id']}"
                           for lv in LEVELS))

    if args.dry_run:
        cameras = pb.load_cameras()
        pools = {lv: pb.load_pool_lines(lv) for lv in LEVELS}
        for item in plan:
            ids = [item["poses"][lv]["id"] for lv in LEVELS]
            jobs = pb.compose_theme_jobs(
                [item["theme"]], 3, dict(QUOTAS), "per-image", {},
                "auto", 0.0, bank, avoid, pools, cameras, rng,
                max_variants=1, bank_ids=ids)
            print(f"\n=== {item['char']} ({item['theme']}) ===")
            for j, job in enumerate(jobs, 1):
                prompt = pb.build_run_prompt(B.master_prompt, item["char_line"], job)
                print(f"[{j}] {job.nivel} pose:{job.pose_id} cam:{job.camera} "
                      f"dress:{job.dress} outfit:{job.variant_keys.get('outfit')}")
                print("    " + prompt.replace("\n", " | ")[-320:])
        return

    if not pb.forge_available():
        raise SystemExit(f"[!] Forge NO responde en {pb.FORGE}. "
                         "Arranca webui-user.bat --api")

    cameras = pb.load_cameras()
    pools = {lv: pb.load_pool_lines(lv) for lv in LEVELS}
    budget = args.budget if args.budget is not None else 3 * 3 * len(plan)

    all_records, plan_out, images = [], [], []
    for n, item in enumerate(plan, 1):
        key = item["char"]
        ids = [item["poses"][lv]["id"] for lv in LEVELS]
        print(f"\n[{n}/{len(plan)}] {key} | tema {item['theme']} | "
              + ", ".join(ids))
        jobs = pb.compose_theme_jobs(
            [item["theme"]], 3, dict(QUOTAS), "per-image", {},
            "auto", 0.0, bank, avoid, pools, cameras, rng,
            max_variants=1, bank_ids=ids)

        got = [j.pose_id for j in jobs]
        if got != ids:
            print(f"[!] aviso: el plan pedia {ids} y salieron {got}")

        char_dir = run_dir / key
        img_dir = char_dir / "img"

        records = load_done(char_dir, ids) if args.resume else None
        if records is not None:
            print(f"    [skip] ya generado ({len(records)} imagenes) -> {char_dir}")
        else:
            records, metrics = pb.generate(jobs, item["char_line"],
                                           B.master_prompt, B.negative_prompt,
                                           img_dir, args.seed, budget, char=key)
            char_dir.mkdir(parents=True, exist_ok=True)
            (char_dir / "pack_manifest.json").write_text(
                json.dumps(records, ensure_ascii=False, indent=1),
                encoding="utf-8")
            (char_dir / "metrics.json").write_text(
                json.dumps(metrics, ensure_ascii=False, indent=1),
                encoding="utf-8")

        for rec in records:
            all_records.append({"char_dir": key, **rec,
                                "file": f"{key}/{rec['file']}"})
            images.append(img_dir / Path(rec["file"]).name)

        plan_out.append({
            "char": key,
            "theme": item["theme"],
            "poses": {lv: {"id": item["poses"][lv]["id"],
                           "line": item["poses"][lv]["pose_line"]}
                      for lv in LEVELS},
            "jobs": [{"nivel": rec["nivel"], "pose_id": rec["pose_id"],
                      "camera": rec["camera"], "dress": rec["dress"],
                      "seed": rec["seed"], "file": rec["file"]}
                     for rec in records],
        })

    run_dir.mkdir(parents=True, exist_ok=True)
    (run_dir / "manifest_all.json").write_text(
        json.dumps(all_records, ensure_ascii=False, indent=1), encoding="utf-8")
    (run_dir / "plan.json").write_text(
        json.dumps({"generated": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                    "chars": len(plan), "seed": args.seed,
                    "pick_seed": args.pick_seed, "plan": plan_out},
                   ensure_ascii=False, indent=1), encoding="utf-8")
    write_contact(images, run_dir / "contact_all.jpg")

    print(f"\n[OK] {len(all_records)}/{len(plan) * 3} imagenes -> {run_dir}")
    print(f"     manifest: {run_dir / 'manifest_all.json'}")
    print(f"     plan:     {run_dir / 'plan.json'}")


if __name__ == "__main__":
    main()