"""
pose_harvest: rescata poses ya renderizadas (origin=explore) y las prepara
para revisión y promoción al banco.

Fase 1 del plan de ampliación: no gasta GPU. Lee los manifiestos de
Output/pilots y Output/packs, limpia las líneas (ropa, accesorios, escena,
atrezo), descarta blacklist y casi-duplicados (Jaccard por conjunto de tags),
y escribe una carpeta de revisión con:

    candidates.json   -> editable: approved true/false, blacklist, motivos
    review.html       -> miniaturas + línea propuesta + origen
    contact_<nivel>.jpg

Promoción (después de revisar y editar candidates.json):

    python Engine/pose_harvest.py --promote Output/pose_review/<stamp>/candidates.json

`approved: true` añade la pose al banco (copia la imagen a refs/). `approved:
false` la descarta sin blacklist. Para blacklist hace falta `blacklist: true`
(y motivos): solo cuando la pose en sí es mala, no cuando la imagen concreta
salió mal.
"""
import argparse
import html
import json
import os
import shutil
import sys
from collections import Counter
from datetime import datetime
from pathlib import Path
from urllib.parse import quote

import pose_bank
from clean_bank_lines import clean_line, tag_key
from wd14_harvest import SCENE

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

BASE_DIR = Path(__file__).resolve().parent.parent
PILOTS_DIR = BASE_DIR / "Output" / "pilots"
PACKS_DIR = BASE_DIR / "Output" / "packs"
REVIEW_ROOT = BASE_DIR / "Output" / "pose_review"

LEVELS = ("sugerente", "spicy", "explicito")
GENERIC = {"1girl", "1boy", "solo", "hetero"}

# Atrezo/acción con objeto que `SCENE` no cubre. La pose de sostener no se
# conserva porque el objeto desaparece y la línea queda coja.
PROP_DROP = {
    "holding_sword", "holding_umbrella", "holding_clothes_hanger",
    "holding_unworn_clothes", "holding_food", "flying", "naked_bandage",
    "bathhouse", "naked_towel",
}


def _content_tokens(line: str) -> set:
    return {t for t in pose_bank.signature(line) if t not in GENERIC}


def _jaccard(a: set, b: set) -> float:
    if not a or not b:
        return 0.0
    return len(a & b) / len(a | b)


def _clean_pose_line(line: str):
    """Quita ropa/meta (clean_bank_lines) y escena/atrezo (SCENE/PROP_DROP)."""
    cleaned, dropped = clean_line(line)
    kept = []
    for token in cleaned.split(","):
        token = token.strip()
        key = tag_key(token)
        if not key:
            continue
        if key in PROP_DROP or SCENE.match(key):
            dropped.append(key)
            continue
        kept.append(token)
    return ", ".join(kept), dropped


def _iter_records(roots):
    for root in roots:
        if not root.exists():
            print(f"[i] no existe {root}; se salta")
            continue
        for manifest in sorted(root.rglob("pack_manifest.json")):
            try:
                data = json.loads(manifest.read_text(encoding="utf-8"))
            except (OSError, json.JSONDecodeError) as e:
                print(f"[!] manifiesto ilegible {manifest}: {e}")
                continue
            records = data.get("records") if isinstance(data, dict) else data
            if not isinstance(records, list):
                continue
            for rec in records:
                if isinstance(rec, dict) and rec.get("origin") == "explore":
                    yield manifest, rec


def _source_for(rec, run_dir: Path):
    rel = rec.get("file") or ""
    if not rel:
        return None
    name = Path(rel).name
    for sub in ("raw", "img"):
        candidate = run_dir / sub / name
        if candidate.exists():
            return candidate.resolve()
    candidate = run_dir / rel
    return candidate.resolve() if candidate.exists() else None


def _bank_index(bank):
    by_level = {}
    for entry in bank:
        by_level.setdefault(entry.get("nivel"), []).append(
            (entry, _content_tokens(entry.get("pose_line", ""))))
    return by_level


def _near_bank(tokens, level, bank_index, threshold):
    best_j, best_entry = 0.0, None
    for entry, sig in bank_index.get(level, []):
        j = _jaccard(tokens, sig)
        if j > best_j:
            best_j, best_entry = j, entry
    return (best_j, best_entry) if best_j >= threshold else (0.0, None)


def collect(roots, min_jaccard=0.7, min_tokens=3):
    bank = pose_bank.load_bank()
    avoid = pose_bank.load_avoid()
    bank_index = _bank_index(bank)
    stats = Counter()
    raw_candidates = []
    seen_sources = set()

    for manifest, rec in _iter_records(roots):
        stats["explore"] += 1
        level = rec.get("nivel")
        if level not in LEVELS:
            stats["nivel_desconocido"] += 1
            continue
        line, dropped = _clean_pose_line(rec.get("pose_line", ""))
        tokens = _content_tokens(line)
        if len(tokens) < min_tokens:
            stats["pocos_tags"] += 1
            continue
        if pose_bank.is_avoided(line, avoid):
            stats["blacklisted"] += 1
            continue
        j, entry = _near_bank(tokens, level, bank_index, min_jaccard)
        if entry is not None:
            stats["casi_duplicado_banco"] += 1
            print(f"[i] {Path(rec.get('file', '')).name}: ~{entry.get('id')} "
                  f"(j={j:.2f})")
            continue
        source = _source_for(rec, manifest.parent)
        if source is None:
            stats["sin_imagen"] += 1
            continue
        key = str(source)
        if key in seen_sources:
            stats["misma_imagen"] += 1
            continue
        seen_sources.add(key)
        raw_candidates.append({
            "nivel": level,
            "pose_line": line,
            "source": os.path.relpath(source, BASE_DIR).replace("\\", "/"),
            "source_name": source.name,
            "run": manifest.parent.name,
            "file": Path(rec.get("file", "")).name,
            "seed": rec.get("seed"),
            "char": rec.get("char", ""),
            "lora_weight": rec.get("lora_weight", 1.0),
            "dropped": sorted(set(dropped)),
            "tokens": tokens,
        })

    # Dedupe entre candidatas: se queda la línea con más tags de contenido.
    kept = []
    for cand in sorted(raw_candidates,
                       key=lambda c: (-len(c["tokens"]), c["source"])):
        duplicate = None
        for other in kept:
            if other["nivel"] != cand["nivel"]:
                continue
            j = _jaccard(cand["tokens"], other["tokens"])
            if j >= min_jaccard:
                duplicate = (j, other)
                break
        if duplicate:
            stats["casi_duplicado_candidato"] += 1
            print(f"[i] {cand['source_name']}: ~{duplicate[1]['source_name']} "
                  f"(j={duplicate[0]:.2f})")
            continue
        kept.append(cand)

    kept.sort(key=lambda c: (LEVELS.index(c["nivel"]), c["source"]))
    for i, cand in enumerate(kept, 1):
        cand["index"] = i
        cand["approved"] = None
        cand["blacklist"] = False
        cand["motivos"] = []
        cand["notas"] = ""
        cand["promoted_id"] = None
        cand.pop("tokens", None)

    by_level = {lv: sum(1 for c in kept if c["nivel"] == lv) for lv in LEVELS}
    summary = {
        "explore_leidos": stats["explore"],
        "candidatas": len(kept),
        "por_nivel": by_level,
        "descartadas": {
            "blacklisted": stats["blacklisted"],
            "casi_duplicado_banco": stats["casi_duplicado_banco"],
            "casi_duplicado_candidato": stats["casi_duplicado_candidato"],
            "pocos_tags": stats["pocos_tags"],
            "sin_imagen": stats["sin_imagen"],
        },
    }
    return kept, summary


def _font(size=18):
    try:
        from PIL import ImageFont
        return ImageFont.truetype("arial.ttf", size)
    except Exception:
        from PIL import ImageFont
        return ImageFont.load_default()


def _contact_sheet(cands, path):
    """Hoja de contacto con el índice global para revisar de un vistazo."""
    from PIL import Image, ImageDraw
    if not cands:
        return False
    cols, tw, th, label_h = 4, 270, 380, 32
    rows = (len(cands) + cols - 1) // cols
    sheet = Image.new("RGB", (cols * tw, rows * (th + label_h)), (15, 21, 32))
    draw = ImageDraw.Draw(sheet)
    font = _font(18)
    for i, cand in enumerate(cands):
        x = (i % cols) * tw
        y = (i // cols) * (th + label_h)
        draw.text((x + 8, y + 6), f"[{cand['index']}] {cand['nivel']}",
                  fill=(230, 237, 243), font=font)
        src = BASE_DIR / cand["source"]
        try:
            img = Image.open(src).convert("RGB")
            img.thumbnail((tw - 12, th - 12))
            sheet.paste(img, (x + (tw - img.width) // 2, y + label_h + 4))
        except Exception as e:
            draw.text((x + 8, y + label_h + 10), f"ERR {e}", fill=(251, 113, 133),
                      font=font)
    sheet.save(path, quality=88)
    return True


def write_review(cands, summary, out_dir, min_jaccard=0.7):
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    payload = {
        "generated": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "min_jaccard": min_jaccard,
        "summary": summary,
        "candidates": cands,
    }
    (out_dir / "candidates.json").write_text(
        json.dumps(payload, ensure_ascii=False, indent=1), encoding="utf-8")

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

    rows = []
    for level in LEVELS:
        subset = [c for c in cands if c["nivel"] == level]
        if not subset:
            continue
        cards = []
        for cand in subset:
            src = (BASE_DIR / cand["source"]).resolve()
            rel = os.path.relpath(src, out_dir).replace("\\", "/")
            rel = quote(rel, safe="/:")
            dropped = ", ".join(cand["dropped"]) or "-"
            cards.append(f"""
      <article class="card">
        <div class="idx">[{cand['index']}]</div>
        <a href="{html.escape(rel)}"><img loading="lazy" src="{html.escape(rel)}"
             alt="{html.escape(cand['source_name'])}"></a>
        <p class="line">{html.escape(cand['pose_line'])}</p>
        <p class="meta">{html.escape(cand['run'])} · s{cand['seed']} ·
           {html.escape(cand['source_name'])}</p>
        <p class="meta">quitado: {html.escape(dropped)}</p>
      </article>""")
        sheet = sheets.get(level)
        link = (f'<a class="sheet" href="{sheet}">hoja de contacto</a>'
                if sheet else "")
        rows.append(f"""
    <section>
      <h2>{level} <span class="count">{len(subset)}</span> {link}</h2>
      <div class="grid">{''.join(cards)}</div>
    </section>""")

    doc = f"""<!DOCTYPE html>
<html lang="es"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Poses · revisión {html.escape(payload['generated'])}</title>
<style>
  :root {{ --bg:#0b0e14; --panel:#121826; --line:#23304a; --text:#e6edf3;
           --muted:#93a4b8; --pink:#ff5c8a; --cyan:#6ee7ff; }}
  * {{ box-sizing:border-box }}
  body {{ margin:0; background:var(--bg); color:var(--text);
          font-family:ui-sans-serif,system-ui,"Segoe UI",Arial,sans-serif; }}
  .wrap {{ max-width:1280px; margin:0 auto; padding:28px 20px 60px; }}
  h1 {{ margin:0 0 6px; font-size:30px; }}
  h2 {{ margin:34px 0 12px; border-bottom:1px solid var(--line);
        padding-bottom:8px; text-transform:capitalize; }}
  .count {{ color:var(--pink); font-size:18px; }}
  .sheet {{ float:right; font-size:13px; color:var(--cyan); font-weight:400; }}
  .sub {{ color:var(--muted); margin:0 0 8px; }}
  .note {{ border-left:3px solid var(--pink); background:#121826;
           padding:12px 14px; border-radius:0 10px 10px 0; color:var(--muted);
           margin:16px 0 8px; }}
  .note code {{ color:var(--text); }}
  .grid {{ display:grid; grid-template-columns:repeat(auto-fill,minmax(260px,1fr));
           gap:14px; }}
  .card {{ position:relative; background:linear-gradient(180deg,#121826,#0f1522);
           border:1px solid var(--line); border-radius:12px; padding:10px;
           overflow:hidden; }}
  .card 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 var(--line); border-radius:999px; padding:2px 9px;
          font-weight:700; color:var(--cyan); font-size:13px; }}
  .line {{ font-size:13px; line-height:1.45; margin:10px 0 8px;
           overflow-wrap:anywhere; }}
  .meta {{ color:var(--muted); font-size:11.5px; margin:2px 0;
           overflow-wrap:anywhere; }}
</style></head><body><div class="wrap">
  <h1>Revisión de poses · {html.escape(payload['generated'])}</h1>
  <p class="sub">{summary['candidatas']} candidatas de
     {summary['explore_leidos']} imágenes explore · Jaccard ≥ {min_jaccard}</p>
  <div class="note">
    Edita <code>candidates.json</code>: <code>approved: true</code> para
    promover, <code>false</code> para descartar. Para blacklist, pon
    <code>blacklist: true</code> y un <code>motivo</code>; solo si la pose es
    mala, no si la imagen concreta salió mal. Después:
    <code>python Engine/pose_harvest.py --promote &lt;esta carpeta&gt;/candidates.json</code>
  </div>
  {''.join(rows)}
</div></body></html>"""
    (out_dir / "review.html").write_text(doc, encoding="utf-8")
    return out_dir


def _source_path(cand):
    src = Path(cand.get("source", ""))
    return src if src.is_absolute() else (BASE_DIR / src)


def _backup_bank():
    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup = pose_bank.BANK_DIR / f"_backup_{stamp}"
    backup.mkdir(parents=True, exist_ok=True)
    if pose_bank.BANK_FILE.exists():
        shutil.copy2(pose_bank.BANK_FILE, backup / pose_bank.BANK_FILE.name)
    if pose_bank.AVOID_FILE.exists():
        shutil.copy2(pose_bank.AVOID_FILE, backup / pose_bank.AVOID_FILE.name)
    return backup


def promote(path, dry_run=False, min_jaccard=None):
    path = Path(path)
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as e:
        print(f"[!] no se puede leer {path}: {e}")
        return 1
    cands = data.get("candidates") if isinstance(data, dict) else None
    if not isinstance(cands, list):
        print("[!] candidates.json sin lista 'candidates'")
        return 1
    if min_jaccard is None:
        min_jaccard = float(data.get("min_jaccard", 0.7))

    bank = pose_bank.load_bank()
    avoid = pose_bank.load_avoid()
    bank_index = _bank_index(bank)
    added_sigs = {lv: [] for lv in LEVELS}

    to_add, to_black, skipped = [], [], []
    for cand in cands:
        if not isinstance(cand, dict):
            continue
        approved = cand.get("approved")
        if approved is None:
            skipped.append((cand, "sin revisar"))
            continue
        line = cand.get("pose_line", "")
        level = cand.get("nivel")
        if approved is False:
            if cand.get("blacklist"):
                to_black.append(cand)
            else:
                skipped.append((cand, "descartada"))
            continue
        if approved is not True:
            skipped.append((cand, "approved no booleano"))
            continue
        if level not in LEVELS:
            skipped.append((cand, "nivel desconocido"))
            continue
        if pose_bank.is_avoided(line, avoid):
            skipped.append((cand, "ya en blacklist"))
            continue
        tokens = _content_tokens(line)
        if len(tokens) < 3:
            skipped.append((cand, "linea demasiado corta"))
            continue
        j, entry = _near_bank(tokens, level, bank_index, min_jaccard)
        if entry is not None:
            skipped.append((cand, f"casi-duplicado de {entry.get('id')}"))
            continue
        if any(_jaccard(tokens, sig) >= min_jaccard
               for sig in added_sigs[level]):
            skipped.append((cand, "casi-duplicado de otra aprobada"))
            continue
        if _source_path(cand).exists() is False:
            skipped.append((cand, "imagen no encontrada"))
            continue
        to_add.append(cand)

    if dry_run:
        print(f"[-] dry-run: {len(to_add)} promover, {len(to_black)} blacklist, "
              f"{len(skipped)} saltadas")
        for cand in to_add:
            print(f"  + [{cand.get('index')}] {cand.get('nivel')}: "
                  f"{cand.get('pose_line')[:100]}")
        for cand, motivo in skipped:
            print(f"  - [{cand.get('index')}] {motivo}")
        return 0

    if not to_add and not to_black:
        print("[i] nada que promover")
        return 0

    backup = _backup_bank()
    print(f"[-] backup del banco: {backup}")

    for cand in to_add:
        src = _source_path(cand)
        try:
            pid = pose_bank.add_pose(
                pose_line=cand.get("pose_line", ""),
                nivel=cand.get("nivel"),
                tema=[],
                ref_png=src,
                char=cand.get("char", ""),
                seed=cand.get("seed"),
                provenance="explore",
                lora_weight=cand.get("lora_weight", 1.0),
                notas=f"promovida de {cand.get('run', '')}",
            )
        except Exception as e:
            print(f"[!] [{cand.get('index')}] no promovida: {e}")
            cand["status"] = f"error: {e}"
            continue
        cand["promoted_id"] = pid
        cand["status"] = "promoted"
        added_sigs[cand["nivel"]].append(_content_tokens(cand["pose_line"]))
        print(f"[+] [{cand.get('index')}] {pid} ({cand.get('nivel')}): "
              f"{cand.get('pose_line')[:90]}")

    for cand in to_black:
        line = cand.get("pose_line", "")
        if pose_bank.is_avoided(line):
            cand["status"] = "ya en blacklist"
            continue
        motivos = cand.get("motivos") or ["revision manual"]
        if isinstance(motivos, str):
            motivos = [motivos]
        pose_bank.penalize(line, ", ".join(motivos), threshold=1)
        cand["status"] = "blacklisted"
        print(f"[x] [{cand.get('index')}] blacklist: {line[:90]}")

    data["candidates"] = cands
    path.write_text(json.dumps(data, ensure_ascii=False, indent=1),
                    encoding="utf-8")
    print(f"\n[OK] promovidas {len(to_add)}, blacklist {len(to_black)}, "
          f"saltadas {len(skipped)}")
    print("[i] revisa con: python Engine/pose_bank.py --stats")
    return 0


def main():
    ap = argparse.ArgumentParser(description="Promueve poses explore al banco")
    ap.add_argument("--promote", metavar="CANDIDATES_JSON",
                    help="promueve las candidatas aprobadas de candidates.json")
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--min-jaccard", type=float, default=0.7,
                    help="umbral de casi-duplicado (default 0.7)")
    ap.add_argument("--min-tokens", type=int, default=3,
                    help="tags de contenido minimos por linea (default 3)")
    ap.add_argument("--roots", nargs="*", default=None,
                    help="carpetas a escanear (default Output/pilots Output/packs)")
    ap.add_argument("--out", default=None,
                    help="carpeta de revision (default Output/pose_review/<stamp>)")
    args = ap.parse_args()

    if args.promote:
        return promote(args.promote, dry_run=args.dry_run,
                       min_jaccard=args.min_jaccard)

    roots = ([Path(p) for p in args.roots] if args.roots
             else [PILOTS_DIR, PACKS_DIR])
    cands, summary = collect(roots, min_jaccard=args.min_jaccard,
                             min_tokens=args.min_tokens)
    out_dir = (Path(args.out) if args.out
               else REVIEW_ROOT / datetime.now().strftime("%Y%m%d_%H%M%S"))
    write_review(cands, summary, out_dir, min_jaccard=args.min_jaccard)

    print(f"[OK] {summary['candidatas']} candidatas "
          f"({summary['por_nivel']}) -> {out_dir}")
    print(f"     descartes: {summary['descartadas']}")
    print(f"     revisa {out_dir / 'review.html'} y edita candidates.json")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())