"""
pose_calibrate: renderiza poses candidatas para revisarlas antes de promoverlas.

Fase 2 del plan de ampliacion: cada candidata se genera con varias seeds y un
fondo neutro para juzgar la pose, no el tema. La salida es compatible con
`pose_harvest.py --promote`:

    python Engine/pose_calibrate.py --dry-run
    python Engine/pose_calibrate.py
    python Engine/pose_harvest.py --promote Output/pose_calibrate/<stamp>/candidates.json

Las candidatas viven en Archivos/pose_candidates.jsonl (una por linea).
"""
import argparse
import html
import json
import sys
from datetime import datetime
from pathlib import Path

import requests

import pose_harvest
from prompt_builder import PromptBuilder
from test_blocks import build_prompt

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

BASE_DIR = Path(__file__).resolve().parent.parent
CANDIDATES = BASE_DIR / "Archivos" / "pose_candidates.jsonl"
OUT_ROOT = BASE_DIR / "Output" / "pose_calibrate"
FORGE = "http://127.0.0.1:7860"
LEVELS = ("sugerente", "spicy", "explicito")
CALIBRATION_EXTRA = "simple_background, white_background"


def load_candidates(path=CANDIDATES):
    out = []
    for n, line in enumerate(Path(path).read_text(encoding="utf-8").splitlines(), 1):
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        rec = json.loads(line)
        if rec.get("nivel") not in LEVELS:
            raise SystemExit(f"[!] {path}:{n}: nivel invalido '{rec.get('nivel')}'")
        if not rec.get("pose_line"):
            raise SystemExit(f"[!] {path}:{n}: falta pose_line")
        rec["id"] = rec.get("id") or f"cand{n:02d}"
        out.append(rec)
    return out


def forge_ready():
    try:
        r = requests.get(f"{FORGE}/sdapi/v1/options", timeout=10)
        return r.status_code == 200, r.json().get("sd_model_checkpoint", "?")
    except Exception as e:
        return False, str(e)


def render(prompt, negative, seed):
    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,
        "restore_faces": False,
        "tiling": False,
        "enable_hr": False,
    }
    r = requests.post(f"{FORGE}/sdapi/v1/txt2img", json=payload, timeout=240)
    r.raise_for_status()
    data = r.json()
    info = data.get("info", "{}")
    if isinstance(info, str):
        try:
            info = json.loads(info)
        except json.JSONDecodeError:
            info = {}
    return data["images"][0], info.get("seed", seed)


def build_candidate(cand, index, source, seed, prompt, char):
    return {
        "nivel": cand["nivel"],
        "pose_line": cand["pose_line"],
        "source": source.relative_to(BASE_DIR).as_posix(),
        "source_name": source.name,
        "run": "pose_calibrate",
        "file": source.name,
        "seed": seed,
        "char": char,
        "dropped": [],
        "index": index,
        "approved": False,
        "blacklist": False,
        "motivos": [],
        "notas": cand.get("notes", ""),
        "candidate_id": cand["id"],
        "prompt": prompt,
    }


def write_review(cands, out_dir, summary):
    sheets = {}
    for level in LEVELS:
        subset = [c for c in cands if c["nivel"] == level]
        if subset and pose_harvest._contact_sheet(subset, out_dir / f"contact_{level}.jpg"):
            sheets[level] = f"contact_{level}.jpg"

    sections = []
    for level in LEVELS:
        subset = [c for c in cands if c["nivel"] == level]
        if not subset:
            continue
        cards = []
        for c in subset:
            cards.append(
                f'<article><div class="idx">[{c["index"]}]</div>'
                f'<a href="{html.escape(c["source"])}"><img loading="lazy" '
                f'src="{html.escape(c["source"])}" alt="{html.escape(c["source_name"])}"></a>'
                f'<p>{html.escape(c["pose_line"])}</p>'
                f'<small>{html.escape(c["candidate_id"])} · seed {c["seed"]}</small></article>'
            )
        link = (f'<a class="sheet" href="{sheets[level]}">hoja de contacto</a>'
                if level in sheets else "")
        sections.append(
            f'<h2>{level} ({len(subset)}) {link}</h2><div class="grid">'
            + "".join(cards) + "</div>"
        )

    doc = f"""<!DOCTYPE html>
<html lang="es"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Poses · calibración {html.escape(summary["generated"])}</title>
<style>
 body{{margin:0;background:#0b0e14;color:#e6edf3;font:15px system-ui,Segoe UI,Arial}}
 .wrap{{max-width:1280px;margin:0 auto;padding:28px 20px 60px}}
 h1{{margin:0 0 6px}} h2{{margin:30px 0 12px;border-bottom:1px solid #23304a;padding-bottom:8px;text-transform:capitalize}}
 .sheet{{float:right;color:#6ee7ff;font-size:13px;font-weight:400}}
 .grid{{display:grid;grid-template-columns:repeat(auto-fill,minmax(250px,1fr));gap:14px}}
 article{{background:#121826;border:1px solid #23304a;border-radius:12px;padding:10px;position:relative}}
 img{{width:100%;height:330px;object-fit:contain;background:#0a0f18;border-radius:8px}}
 .idx{{position:absolute;top:16px;left:16px;background:rgba(11,14,20,.82);border:1px solid #23304a;border-radius:999px;padding:2px 9px;color:#6ee7ff;font-weight:700}}
 p{{font-size:13px;line-height:1.4;overflow-wrap:anywhere}} small{{color:#93a4b8}}
</style></head><body><div class="wrap">
<h1>Calibración de poses · {html.escape(summary["generated"])}</h1>
<p>{summary["poses"]} poses candidatas × {summary["seeds"]} seeds = {summary["images"]} imágenes.
Edita <code>candidates.json</code>: deja <code>approved: true</code> en la mejor imagen
de cada pose y descarta el resto. Después:
<code>python Engine/pose_harvest.py --promote &lt;esta carpeta&gt;/candidates.json</code></p>
{''.join(sections)}
</div></body></html>"""
    (out_dir / "review.html").write_text(doc, encoding="utf-8")


def main():
    p = argparse.ArgumentParser(description="Renderiza poses candidatas para calibrar")
    p.add_argument("--dry-run", action="store_true", help="solo imprime prompts")
    p.add_argument("--char", default="Makima_v2")
    p.add_argument("--seeds", default="42,123,777")
    p.add_argument("--only", default=None, help="ids de candidata, separados por coma")
    p.add_argument("--candidates", default=str(CANDIDATES),
                   help="jsonl de candidatas (default Archivos/pose_candidates.jsonl)")
    p.add_argument("--one-seed-per-pose", action="store_true",
                   help="una semilla DISTINTA por pose: manda -1 (aleatoria) y "
                        "guarda la que devuelve Forge en cada imagen")
    p.add_argument("--seed-base", type=int, default=None,
                   help="opcional: base fija (seed-base + i*7919) para repetir "
                        "un lote concreto; sin esto se usa -1 (aleatoria)")
    p.add_argument("--out", default=None)
    args = p.parse_args()

    cands = load_candidates(args.candidates)
    if args.only:
        wanted = {x.strip() for x in args.only.split(",") if x.strip()}
        cands = [c for c in cands if c["id"] in wanted]
    if args.one_seed_per_pose:
        # -1 = Forge tira una semilla aleatoria distinta en cada generacion y la
        # devuelve en `info.seed`; esa es la que se guarda en el manifest.
        # Con --seed-base se puede reproducir un lote concreto.
        if args.seed_base is None:
            pairs = [(c, -1) for c in cands]
        else:
            pairs = [(c, args.seed_base + i * 7919) for i, c in enumerate(cands, 1)]
        seeds = [s for _, s in pairs]
    else:
        seeds = [int(x.strip()) for x in args.seeds.split(",") if x.strip()]
        pairs = [(c, s) for c in cands for s in seeds]
    if not cands or not seeds:
        raise SystemExit("[!] no hay candidatas o seeds")

    B = PromptBuilder(str(BASE_DIR))
    char = B.get_character(args.char)
    if not char:
        raise SystemExit(f"[!] personaje '{args.char}' no encontrado")
    master = B.master_prompt
    negative = B.negative_prompt
    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    out_dir = Path(args.out) if args.out else OUT_ROOT / stamp
    img_dir = out_dir / "img"

    if args.one_seed_per_pose:
        modo = "aleatoria por pose (-1)" if args.seed_base is None else \
            f"fija por pose ({seeds[0]}..{seeds[-1]})"
        print(f"[-] {len(pairs)} poses, semilla {modo}")
    else:
        print(f"[-] {len(cands)} poses x {len(seeds)} seeds = {len(pairs)} imagenes")
    print(f"[-] personaje: {args.char} | salida: {out_dir}")

    if args.dry_run:
        for c, seed in pairs:
            block = f'{c["pose_line"]}, {CALIBRATION_EXTRA}'
            prompt = build_prompt(master, char["base"], block, force_nude="always")
            print(f"\n[{c['id']}] {c['nivel']} seed={seed} | {c['pose_line']}")
            print(f"    {prompt[-500:].replace(chr(10), ' | ')}")
        return

    ok, ckpt = forge_ready()
    if not ok:
        raise SystemExit(f"[!] Forge no responde en {FORGE}: {ckpt}")
    print(f"[-] Forge OK | checkpoint: {ckpt}")

    img_dir.mkdir(parents=True, exist_ok=True)
    records = []
    manifest = []
    index = 0
    for c, seed in pairs:
        block = f'{c["pose_line"]}, {CALIBRATION_EXTRA}'
        prompt = build_prompt(master, char["base"], block, force_nude="always")
        index += 1
        print(f"[{index:02d}] {c['id']} seed={seed}", end=" ... ", flush=True)
        try:
            b64, used_seed = render(prompt, negative, seed)
            import base64
            # el nombre lleva el id de la pose (y la semilla REAL que devolvio
            # Forge, no el -1 de la peticion)
            name = f"{index:02d}_{c['nivel']}_{c['id']}_s{used_seed}.png"
            (img_dir / name).write_bytes(base64.b64decode(b64))
        except Exception as e:
            print(f"FALLO: {e}")
            continue
        source = img_dir / name
        records.append(build_candidate(c, index, source, used_seed, prompt, args.char))
        manifest.append({
            "file": source.relative_to(BASE_DIR).as_posix(),
            "candidate_id": c["id"],
            "nivel": c["nivel"],
            "seed": used_seed,
            "prompt": prompt,
            "negative": negative,
        })
        print(f"OK s{used_seed}")

    out_dir.mkdir(parents=True, exist_ok=True)
    payload = {
        "generated": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "char": args.char,
        "seeds": seeds,
        "mode": "one-seed-per-pose" if args.one_seed_per_pose else "seeds",
        "summary": {"poses": len(cands),
                    "seeds": 1 if args.one_seed_per_pose else len(seeds),
                    "images": len(records)},
        "candidates": records,
    }
    payload["summary"]["generated"] = payload["generated"]
    (out_dir / "candidates.json").write_text(
        json.dumps(payload, ensure_ascii=False, indent=1), encoding="utf-8")
    (out_dir / "manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=1), encoding="utf-8")
    write_review(records, out_dir, payload["summary"])
    print(f"[OK] {len(records)} imagenes -> {out_dir}")
    print(f"     revisa {out_dir / 'review.html'} y edita candidates.json")


if __name__ == "__main__":
    main()