"""
pose_bank: banco de poses validadas (JSONL) + CLI. Es la base compartida del
pipeline de contenido (pack_builder lo consume, la auditoria lo promociona y
penaliza).

Ruta por defecto: Archivos/pose_bank/
    pose_bank.jsonl -> una pose validada por linea
    avoid.jsonl     -> poses/bloques fallidos (con blacklist automatica)
    refs/<id>.png   -> render de referencia aprobado

Se puede apuntar a otro directorio con la variable de entorno PLUMLOOM_BANK_DIR
(lo usan las verificaciones para no tocar el banco real).

    python Engine/pose_bank.py --stats
    python Engine/pose_bank.py --list
    python Engine/pose_bank.py --check
    python Engine/pose_bank.py --avoid KEY --motivo M [--threshold N]
    python Engine/pose_bank.py --unblacklist KEY
    python Engine/pose_bank.py --promote RUN_DIR --verdicts PATH
"""
import argparse
import json
import logging
import os
import random
import re
import shutil
import sys
from contextlib import contextmanager
from datetime import datetime
from pathlib import Path

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

BASE_DIR = Path(__file__).resolve().parent.parent
logger = logging.getLogger(__name__)

BANK_DIR = (Path(os.environ["PLUMLOOM_BANK_DIR"])
            if os.environ.get("PLUMLOOM_BANK_DIR")
            else BASE_DIR / "Archivos" / "pose_bank")
BANK_FILE = BANK_DIR / "pose_bank.jsonl"
AVOID_FILE = BANK_DIR / "avoid.jsonl"
REFS_DIR = BANK_DIR / "refs"

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

_WEIGHT_RE = re.compile(r"\(\s*([^():]*?)\s*:\s*[\d.]+\s*\)")
_ID_RE = re.compile(r"^pb(\d+)$")


def ensure_dirs() -> None:
    BANK_DIR.mkdir(parents=True, exist_ok=True)
    REFS_DIR.mkdir(parents=True, exist_ok=True)


def _read_jsonl(path: Path) -> list[dict]:
    if not path.exists():
        return []
    out = []
    with path.open("r", encoding="utf-8") as f:
        for lineno, line in enumerate(f, 1):
            if not line.strip():
                continue
            try:
                out.append(json.loads(line))
            except json.JSONDecodeError as e:
                logger.warning("%s:%d linea corrupta, saltada (%s)", path.name, lineno, e)
    return out


def _write_jsonl(path: Path, rows: list[dict]) -> None:
    ensure_dirs()
    with path.open("w", encoding="utf-8", newline="\n") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")


def load_bank() -> list[dict]:
    return _read_jsonl(BANK_FILE)


def save_bank(entries: list[dict]) -> None:
    _write_jsonl(BANK_FILE, entries)


def load_avoid() -> dict[str, dict]:
    avoid = {}
    for row in _read_jsonl(AVOID_FILE):
        key = row.get("key")
        if key:
            avoid[key] = row
    return avoid


def save_avoid(d: dict) -> None:
    _write_jsonl(AVOID_FILE, list(d.values()))


def normalize_key(pose_line: str) -> str:
    s = (pose_line or "").lower()
    s = _WEIGHT_RE.sub(r"\1", s)
    s = s.replace("_", " ")
    return re.sub(r"\s+", " ", s).strip()


def signature(pose_line: str) -> tuple[str, ...]:
    """Tags normalizados y ordenados, sin duplicados.

    `normalize_key` conserva el orden (lo usa avoid.jsonl, que no se puede
    romper). Para detectar casi-duplicados hace falta una firma insensible al
    orden: dos lineas con los mismos tags en distinto orden son la misma pose.
    """
    tags = {
        normalize_key(token)
        for token in (pose_line or "").split(",")
        if normalize_key(token)
    }
    return tuple(sorted(tags))


def next_id(entries) -> str:
    top = 0
    for e in entries:
        m = _ID_RE.match(str(e.get("id", "")))
        if m:
            top = max(top, int(m.group(1)))
    return f"pb{top + 1:04d}"


def add_pose(pose_line, nivel, tema, ref_png, char, seed, provenance,
             lora_weight=1.0, notas="") -> str:
    entries = load_bank()
    new_id = next_id(entries)

    if isinstance(tema, str):
        tema = [tema]
    tema = list(tema or [])

    ref_stored = None
    if ref_png:
        src = Path(ref_png)
        if src.exists():
            ensure_dirs()
            dest = REFS_DIR / f"{new_id}.png"
            shutil.copy2(src, dest)
            ref_stored = f"refs/{new_id}.png"
        else:
            ref_stored = ref_png

    entry = {
        "id": new_id,
        "pose_line": pose_line,
        "nivel": nivel,
        "tema": tema,
        "provenance": provenance,
        "ref_png": ref_stored,
        "char": char,
        "seed": seed,
        "lora_weight": lora_weight,
        "fecha": datetime.now().strftime("%Y-%m-%d"),
        "notas": notas,
    }
    entries.append(entry)
    save_bank(entries)
    return new_id


def is_avoided(pose_line, avoid=None) -> bool:
    if avoid is None:
        avoid = load_avoid()
    row = avoid.get(normalize_key(pose_line))
    return bool(row and row.get("blacklist"))


def blacklist(pose_line, motivo) -> dict:
    """Mete la pose en avoid con blacklist directa (la usa la UI, D2.3).

    A diferencia de `penalize`, no cuenta fallos: deja la pose fuera del pool
    ya. Es idempotente (volver a llamar no cambia nada).
    """
    key = normalize_key(pose_line)
    avoid = load_avoid()
    row = avoid.get(key) or {
        "key": key, "motivo": motivo, "fallos": 0, "blacklist": False, "ultima": "",
    }
    row["motivo"] = motivo
    row["blacklist"] = True
    row["ultima"] = datetime.now().strftime("%Y-%m-%d")
    avoid[key] = row
    save_avoid(avoid)
    return row


def unblacklist(pose_line) -> bool:
    """Saca la pose de avoid (reactivar). -> True si estaba."""
    key = normalize_key(pose_line)
    avoid = load_avoid()
    if key not in avoid:
        return False
    del avoid[key]
    save_avoid(avoid)
    return True


def penalize(pose_line, motivo, threshold=3) -> dict:
    key = normalize_key(pose_line)
    avoid = load_avoid()
    row = avoid.get(key) or {
        "key": key, "motivo": motivo, "fallos": 0, "blacklist": False, "ultima": "",
    }
    row["motivo"] = motivo
    row["fallos"] = int(row.get("fallos", 0)) + 1
    row["ultima"] = datetime.now().strftime("%Y-%m-%d")
    row["blacklist"] = row["fallos"] >= threshold
    avoid[key] = row
    save_avoid(avoid)
    return row


@contextmanager
def _bank_dir_override(bank_dir):
    """Apunta temporalmente el banco a `bank_dir` (None = el del modulo)."""
    global BANK_DIR, BANK_FILE, AVOID_FILE, REFS_DIR
    if bank_dir is None:
        yield
        return
    saved = (BANK_DIR, BANK_FILE, AVOID_FILE, REFS_DIR)
    BANK_DIR = Path(bank_dir)
    BANK_FILE = BANK_DIR / "pose_bank.jsonl"
    AVOID_FILE = BANK_DIR / "avoid.jsonl"
    REFS_DIR = BANK_DIR / "refs"
    try:
        yield
    finally:
        BANK_DIR, BANK_FILE, AVOID_FILE, REFS_DIR = saved


def load_verdicts(path) -> dict:
    """Lee un verdicts_<stamp>.json (dict por nombre de fichero).

    Devuelve {} (sin lanzar) si el fichero falta, no es JSON o no es un dict.
    """
    path = Path(path)
    if not path.exists():
        logger.warning("verdicts no encontrado: %s", path)
        return {}
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, ValueError, OSError) as e:
        logger.warning("verdicts ilegible %s (%s)", path, e)
        return {}
    if not isinstance(data, dict):
        logger.warning("verdicts con forma inesperada en %s (dict esperado)", path)
        return {}
    return data


def promote_run(run_dir, verdicts_path, bank_dir=None) -> dict:
    """Pliega los veredictos de una corrida en el banco.

    - `aprobada` + `origin=="explore"` -> `add_pose` (ref_png = raw/<file> si
      existe, si no img/<file>; original limpio preferido).
    - `rechazada` -> `penalize` con `threshold=1`: un rechazo de auditoria
      excluye la pose en TODAS partes de inmediato (Review Focus #5), ya que la
      clave avoid es la linea de pose normalizada. Idempotente: si la pose ya
      esta en avoid, se salta (no re-incrementa `fallos`).
    - `dudosa` -> skipped (ni promociona ni penaliza).
    - `origin=="bank"` aprobada -> skipped (ya estaba en el banco).
    Nunca duplica: si `normalize_key(pose_line)` ya esta en el banco para ese
    nivel, se salta. Escribe `verdict`/`motivos` de cada registro de vuelta en
    `pack_manifest.json`. -> {"promoted":int,"penalized":int,"skipped":int}
    """
    run_dir = Path(run_dir)
    result = {"promoted": 0, "penalized": 0, "skipped": 0}

    manifest_path = run_dir / "pack_manifest.json"
    if not manifest_path.exists():
        print(f"[!] sin pack_manifest.json en {run_dir}; nada que promover")
        return result

    verdicts = load_verdicts(verdicts_path)
    if not verdicts:
        print(f"[!] sin verdicts en {verdicts_path}; nada que promover")
        return result

    try:
        records = json.loads(manifest_path.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, ValueError, OSError) as e:
        print(f"[!] manifest ilegible en {manifest_path}: {e}")
        return result
    if not isinstance(records, list):
        print(f"[!] manifest con forma inesperada en {manifest_path}")
        return result

    with _bank_dir_override(bank_dir):
        bank = load_bank()
        known = {(normalize_key(e.get("pose_line", "")), e.get("nivel"))
                 for e in bank}

        for rec in records:
            if not isinstance(rec, dict):
                continue
            name = Path(rec.get("file", "")).name
            v = verdicts.get(name)
            if v is None:
                v = verdicts.get(rec.get("file"))
            if not isinstance(v, dict):
                result["skipped"] += 1
                continue

            verdict = str(v.get("verdict") or "").strip().lower()
            motivos = list(v.get("motivos") or [])
            rec["verdict"] = verdict or None
            rec["motivos"] = motivos

            if verdict == "aprobada":
                if rec.get("origin") != "explore":
                    result["skipped"] += 1
                    continue
                pose_line = rec.get("pose_line", "")
                key = (normalize_key(pose_line), rec.get("nivel"))
                if key in known:
                    result["skipped"] += 1
                    continue
                ref = run_dir / "raw" / name
                if not ref.exists():
                    ref = run_dir / "img" / name
                add_pose(pose_line=pose_line, nivel=rec.get("nivel"),
                         tema=rec.get("tema") or [],
                         ref_png=ref if ref.exists() else None,
                         char=rec.get("char", ""), seed=rec.get("seed"),
                         provenance="explore",
                         lora_weight=rec.get("lora_weight", 1.0))
                known.add(key)
                result["promoted"] += 1
            elif verdict == "rechazada":
                pose_line = rec.get("pose_line", "")
                if is_avoided(pose_line):
                    # ya penalizada (por esta corrida o una anterior): no
                    # volver a incrementar `fallos` (idempotente por corrida).
                    result["skipped"] += 1
                    continue
                penalize(pose_line, ", ".join(motivos), threshold=1)
                result["penalized"] += 1
            else:
                result["skipped"] += 1

    manifest_path.write_text(
        json.dumps(records, ensure_ascii=False, indent=1), encoding="utf-8")
    return result


def sample(entries, nivel=None, tema=None, exclude_ids=None, n=1, rng=None) -> list[dict]:
    rng = rng or random
    exclude = set(exclude_ids or [])
    avoid = load_avoid()

    pool = entries
    if nivel is not None:
        pool = [e for e in pool if e.get("nivel") == nivel]
    if tema is not None:
        tagged = [e for e in pool if tema in (e.get("tema") or [])]
        if tagged:
            pool = tagged

    pool = [e for e in pool
            if e.get("id") not in exclude
            and not is_avoided(e.get("pose_line", ""), avoid)]

    n = max(0, int(n))
    if n <= 0:
        return []
    if n >= len(pool):
        chosen = list(pool)
        rng.shuffle(chosen)
        return chosen
    return rng.sample(pool, n)


def stats(entries=None) -> dict:
    if entries is None:
        entries = load_bank()
    by_level = {lv: 0 for lv in LEVELS}
    temas = set()
    for e in entries:
        if e.get("nivel") in by_level:
            by_level[e["nivel"]] += 1
        temas.update(e.get("tema") or [])
    avoided = sum(1 for row in load_avoid().values() if row.get("blacklist"))
    return {"total": len(entries), "by_level": by_level,
            "temas": sorted(temas), "avoided": avoided}


def check_integrity(entries=None) -> list[str]:
    if entries is None:
        entries = load_bank()
    problems = []
    seen = set()
    for e in entries:
        eid = e.get("id")
        if eid in seen:
            problems.append(f"duplicate id: {eid}")
        seen.add(eid)
        ref = e.get("ref_png")
        if ref:
            p = Path(ref)
            if not p.is_absolute():
                p = BANK_DIR / ref
            if not p.exists():
                problems.append(f"{eid}: ref_png missing on disk: {ref}")
    return problems


def main():
    logging.basicConfig(level=logging.WARNING, format="%(levelname)s %(message)s")
    p = argparse.ArgumentParser(description="Banco de poses PlumLoom (pose_bank)")
    p.add_argument("--list", action="store_true", help="lista las poses del banco")
    p.add_argument("--stats", action="store_true", help="resumen por nivel/tema")
    p.add_argument("--check", action="store_true", help="chequeo de integridad")
    p.add_argument("--avoid", metavar="KEY", help="penaliza una pose")
    p.add_argument("--motivo", help="motivo de la penalizacion (con --avoid)")
    p.add_argument("--threshold", type=int, default=3,
                   help="fallos para blacklist automatica (default 3)")
    p.add_argument("--unblacklist", metavar="KEY", help="saca una pose del avoid")
    p.add_argument("--promote", metavar="RUN_DIR",
                   help="promociona los aprobados de una corrida (con --verdicts)")
    p.add_argument("--verdicts", metavar="PATH",
                   help="verdicts_<stamp>.json (con --promote)")
    args = p.parse_args()

    if args.promote:
        if not args.verdicts:
            raise SystemExit("[!] --promote requiere --verdicts")
        res = promote_run(args.promote, args.verdicts)
        print(f"promoted {res['promoted']} penalized {res['penalized']} "
              f"skipped {res['skipped']}")
        return

    if args.avoid:
        if not args.motivo:
            raise SystemExit("[!] --avoid requiere --motivo")
        row = penalize(args.avoid, args.motivo, args.threshold)
        print(json.dumps(row, ensure_ascii=False))
        return

    if args.unblacklist:
        avoid = load_avoid()
        key = normalize_key(args.unblacklist)
        if key in avoid:
            del avoid[key]
            save_avoid(avoid)
            print(f"unblacklisted: {key}")
        else:
            print(f"no estaba en avoid: {key}")
        return

    if args.check:
        problems = check_integrity()
        for pr in problems:
            print(f"[!] {pr}")
        if problems:
            raise SystemExit(1)
        print("[OK] banco integro")
        return

    if args.stats:
        s = stats()
        print(f"total: {s['total']}")
        print("by_level: " + ", ".join(f"{k}={v}" for k, v in s["by_level"].items()))
        print("temas: " + (", ".join(s["temas"]) if s["temas"] else "-"))
        print(f"avoided: {s['avoided']}")
        return

    entries = load_bank()
    for e in entries:
        temas = ",".join(e.get("tema") or [])
        print(f"{e.get('id')}  [{e.get('nivel')}]  {temas}  {e.get('pose_line', '')}")
    print(f"({len(entries)} poses en {BANK_FILE})")


if __name__ == "__main__":
    main()