"""
Test de BLOQUES: mette un bloque curado de wd14_blocks.py tras un personaje
del catalogo y comprueba si Stable Diffusion respeta la pose.

La diferencia con test_harvest_tags.py: aqui NO se recombinan pose/outfit/scene
por separado. Se coge el bloque entero tal cual (una linea de bloques.txt) y se
pega DETRAS de la linea del personaje. Asi se mide lo que el bloque pide, no lo
que nosotros reconstruimos.

    # 1. Ver los prompts sin gastar GPU
    python Engine/test_blocks.py --dry-run

    # 2. Un bloque concreto, 2 variaciones
    python Engine/test_blocks.py --char Seed --block 0 --variations 2

    # 3. Barrido: 6 bloques sobre varios personajes
    python Engine/test_blocks.py --n 6

    # 4. A/B focalizado de ropa: mismo indice en ambos pools
    python Engine/test_blocks.py --block 5 --pool desnudo
    python Engine/test_blocks.py --block 5 --pool vestido

Salida en Output/test_blocks/ + manifest_*.json con el prompt exacto de cada
imagen, y una copia del .txt original para comparar lado a lado.
"""
import argparse
import base64
import json
import random
import re
import sys
import time
from datetime import datetime
from pathlib import Path

BASE_DIR = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(Path(__file__).resolve().parent))
from prompt_builder import PromptBuilder  # noqa: E402

BLOCKS = BASE_DIR / "Archivos" / "wildcards_harvest" / "bloques.txt"
BLOCKS_DRESSED = BASE_DIR / "Archivos" / "wildcards_harvest" / "bloques_vestidos.txt"
TAGS_DIR = BASE_DIR / "Archivos" / "wd14" / "tags"
OUT_ROOT = BASE_DIR / "Output" / "test_blocks"
FORGE = "http://127.0.0.1:7860"

# Una carpeta por pool: Output/test_blocks/sfw/, /spicy/, /explicit/...
# Asi las revisiones por severidad no se mezclan (pedido 07-10-2026).
def out_dir(pool):
    return OUT_ROOT / pool

HARVEST = BASE_DIR / "Archivos" / "wildcards_harvest"
HARVEST_DAN = BASE_DIR / "Archivos" / "wildcards_danbooru"
sys.path.insert(0, str(Path(__file__).resolve().parent))
from wd14_blocks import strip_weights_of  # noqa: E402

# Los personajes de personajes.txt son BASE: ya salen desnudos, asi que rotar
# vestido/desnudo con ellos es inutil (medido). El A/B de ropa solo tiene
# sentido con el pool vestido, donde el bloque trae la ropa del WD14 sin color.
POOLS = {
    "desnudo": BLOCKS,
    "vestido": BLOCKS_DRESSED,
    "sfw": HARVEST / "bloques_sfw.txt",
    "spicy": HARVEST / "bloques_spicy.txt",
    "explicit": HARVEST / "bloques_explicit.txt",
"sfw_vestido": HARVEST / "bloques_sfw_vestidos.txt",
              "spicy_vestido": HARVEST / "bloques_spicy_vestidos.txt",
              "explicit_vestido": HARVEST / "bloques_explicit_vestidos.txt",
              # muestra de Danbooru (tags de la comunidad, no de WD14)
              "dan_sfw": HARVEST_DAN / "bloques_sfw.txt",
              "dan_spicy": HARVEST_DAN / "bloques_spicy.txt",
              "dan_explicit": HARVEST_DAN / "bloques_explicit.txt",
              "dan_todos": HARVEST_DAN / "bloques.txt",
              }

# El hombre siempre faceless (ver AGENTS.md / Config/male_template.txt). Un
# bloque con 1boy + pose de dos personas sin esto produce una caraDeformada.
MALE_TAGS = (
    "1boy, (faceless male:1.4), no face, blank face, featureless face, "
    "no facial features, eyes closed, black hair, short hair, "
    "muscular body, broad shoulders, large hands, masculine physique, "
    "nude male, larger male, taller dominant"
)

# Candidatos. OJO: la clave de PromptBuilder es el nombre del LoRA, no el
# nombre del personaje (ver AGENTS.md). Seed = Hoseki_ZenlessZoneZero_Seed_...
CHARS = [
    "Hoseki_ZenlessZoneZero_Seed_IllustriousXL_v1",
    "Arlecchino", "Yelan", "rem", "ReZero_Rem_IlluXL", "nilou-illu-nvwls-v1",
    "Furina", "ai-hoshino-s1-illustriousxl-lora-nochekaiser",
    "Raiden_shogun", "hu_tao_genshin_impact", "Skirk", "kugisaki_nobara",
    "sanae", "Fischl", "sakamata_chloe", "firefly_honkai_star_rail",
    "robin_hsr", "gold_ship_illust_scarxzys",
]

# Freno de camara. Los bloques con ass/pussy/nipples empujan el encuadre hasta
# que el cuerpo llena el cuadro y se deforma (medido: Raiden b02, Yelan b32).
# Esto obliga a la camara a retroceder. Va DESPUES del bloque para ganar.
CAMERA_BRAKE = "(full body:1.2), wide shot, whole body visible, both characters fully visible"


def load_blocks(pool="desnudo"):
    path = POOLS.get(pool, BLOCKS)
    if not path.exists():
        raise SystemExit(f"[!] No existe {path}. Ejecuta antes wd14_blocks.py")
    out = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if line and not line.startswith("#"):
            out.append(line)
    if not out:
        raise SystemExit(f"[!] {BLOCKS} no tiene bloques")
    return out


def find_original(block, tags_dir=TAGS_DIR):
    """Localiza el .txt de Archivos/wd14/tags del que salio este bloque.

    Compara SIN pesos: bloques.txt lleva '(ass_focus:0.7)' pero el .txt original
    no, asi que hay que normalizar antes de comparar.
    """
    from wd14_blocks import load_character_tags, clean_block
    chars, copies, alias = load_character_tags()
    want = strip_weights(block)
    for f in sorted(tags_dir.rglob("*.txt")):
        raw = [t.strip() for t in f.read_text(encoding="utf-8").split(",") if t.strip()]
        keep, _ = clean_block([alias.get(t.lower(), t.lower()) for t in raw], chars, copies)
        seen, uniq = set(), []
        for t in keep:
            if t not in seen:
                seen.add(t)
                uniq.append(t)
        if ", ".join(uniq) == want:
            return f
    return None


def strip_weights(block):
    """'(ass_focus:0.7), nude' -> 'ass_focus, nude'. Para comparar con el .txt."""
    return re.sub(r"\(([^:()]+):[\d.]+\)", r"\1", block)


def lora_on_disk(name):
    d = Path(r"D:\Data\Packages\Stable Diffusion WebUI Forge - Neo\models\Lora")
    if not d.exists():
        return None
    stems = {p.stem.lower() for p in d.rglob("*.safetensors")}
    if name.lower() in stems:
        return True
    return name.rsplit("-", 1)[0].lower() in stems


def block_has_clothing(block):
    """True si el bloque trae alguna prenda (mismo criterio que wd14_blocks)."""
    from wd14_blocks import DROP_OUTFIT
    return any(DROP_OUTFIT.search(t.lower()) for t in block.split(","))


# --- nuevo ensamblador (07-10-2026) ------------------------------------------------
# Estructura con BREAK, segun decision del usuario:
#   SOLO   : MASTER+LoRAs / 1girl+personaje+nude / BREAK / bloque [BREAK escena]
#   PAREJA : MASTER+LoRAs / 1girl+personaje+nude / BREAK / 1boy+hombre / BREAK /
#            bloque [BREAK escena]
# El hombre se simplifica, y en POV se reduce a '1boy, muscular body' porque
# '(faceless male:1.4)' llegaba a crear un segundo hombre.
MALE_FULL = "1boy, (faceless male:1.4), muscular body, broad shoulders, taller male"
MALE_POV = "1boy, muscular body"
COUNT_TAGS = ("1girl", "1boy")   # maximo una vez: '1boy, 1boy' -> 2 hombres
NUDE_TAGS = "nude, completely_nude, naked"


def assemble_prompt(master, char_line, block, force_nude="auto", age_up=False):
    """-> (prompt_con_BREAK, [tramos])"""
    block_clean = strip_weights_of(block)
    is_pov = "pov" in {t.strip().lower() for t in block_clean.split(",")}
    two_people = ("1boy" in block_clean or "hetero" in block_clean)

    # el MASTER trae '1girl, solo'. Con dos personas 'solo' pelea con la pose
    # y hace que el modelo ignore al hombre.
    if two_people:
        master = re.sub(r"\bsolo\b", "", master)
        master = re.sub(r",\s*,", ",", master).strip().strip(",")

    # tramo 1: la chica + su linea + la desnudez si toca
    girl = ["1girl", char_line]
    if age_up:
        girl.append("age up, mature female")
    if force_nude == "always" or (force_nude == "auto"
                                  and not block_has_clothing(block)):
        girl.append(NUDE_TAGS)

    segs = [master, ", ".join(girl)]

    # tramo 2 (solo si hay hombre): el hombre, que va tras el primer BREAK
    if two_people:
        segs.append(MALE_POV if is_pov else MALE_FULL)

    # tramo 3: el bloque de pose, tal cual
    segs.append(block_clean)

    # dedupe de conteos GLOBAL: '1boy' repetido (inyectado + del bloque) hace
    # que el modelo dibuje dos hombres
    seen = set()
    out = []
    for s in segs:
        toks = [t.strip() for t in s.split(",") if t.strip()]
        kept = []
        for t in toks:
            key = t.lower()
            if key in COUNT_TAGS:
                if key in seen:
                    continue
                seen.add(key)
            kept.append(t)
        out.append(", ".join(kept))
    return "\n\nBREAK\n".join(x for x in out if x.strip()), out


def build_prompt(master, char_line, block, cam_brake=True, age_up=False,
                 force_nude="auto"):
    """Compatibilidad: usa el ensamblador nuevo. cam_brake ya no se aplica (los
    tags de cuerpo completo se descartaron por decision del usuario)."""
    pos, _ = assemble_prompt(master, char_line, block,
                             force_nude=force_nude, age_up=age_up)
    return pos


def main():
    p = argparse.ArgumentParser(description="Test de bloques WD14 sobre personajes")
    p.add_argument("--block", type=int, default=None, help="indice de bloque (0-based)")
    p.add_argument("--char", help="personaje concreto (si no, se elige al azar)")
    p.add_argument("--n", type=int, default=4, help="cuantos bloques probar (al azar)")
    p.add_argument("--all", action="store_true",
                   help="recorre TODOS los bloques del pool en orden 0..N-1, "
                        "para auditarlos uno por uno")
    p.add_argument("--skip", type=int, default=0,
                   help="con --all, empieza por este indice (para reanudar)")
    p.add_argument("--only", type=int, action="append", default=None,
                   help="con --all, procesa solo estos indices (repetible). "
                        "Ej: --only 3 --only 7")
    p.add_argument("--variations", type=int, default=1, help="seeds por bloque")
    p.add_argument("--pool", choices=list(POOLS), default="desnudo",
                   help="desnudo = bloques.txt (sin ropa) | vestido = bloques_vestidos.txt "
                        "(ropa del WD14 sin color). Ambos pools comparten indice.")
    p.add_argument("--seed", type=int, default=-1)
    p.add_argument("--lora-weight", type=float, default=None)
    p.add_argument("--no-cam-brake", action="store_true",
                   help="quita el freno de camara (A/B: ver si deforma menos)")
    p.add_argument("--age-up", action="store_true",
                   help="añade 'age up, mature female'. FUERA por defecto: fuerza "
                        "rasgos maduros que chocan con personajes de busto pequeno")
    p.add_argument("--force-nude", choices=["auto", "always", "never"], default="auto",
                   help="auto (defecto) = fuerza nude solo si el bloque no trae ropa "
                        "(para LoRA con ropa horneada, p.ej. Yelan). "
                        "always = siempre. never = nunca.")
    p.add_argument("--dry-run", action="store_true")
    args = p.parse_args()

    blocks = load_blocks(args.pool)
    B = PromptBuilder(str(BASE_DIR))

    import requests
    try:
        r = requests.get(f"{FORGE}/sdapi/v1/options", timeout=15)
        ckpt = r.json().get("sd_model_checkpoint", "?")
        print(f"[-] Forge OK | checkpoint: {ckpt}")
    except Exception as e:
        print(f"[!] Forge NO responde ({e}). Arranca webui-user.bat --api")

    usable = [c for c in CHARS if B.get_character(c) and lora_on_disk(c) is True]
    if not usable:
        usable = [c for c in CHARS if B.get_character(c)]
    print(f"[-] {len(blocks)} bloques (pool: {args.pool}) | personajes usables: {len(usable)}")
    for c in CHARS:
        if B.get_character(c) and lora_on_disk(c) is not True:
            print(f"    AVISO sin LoRA en disco -> {c}")

    # seleccionar bloques
    if args.all:
        # modo secuencial 0..N-1 para auditar TODOS los bloques uno por uno.
        # --skip permite reanudar donde se quedo la corrida anterior.
        idxs = [i for i in range(len(blocks)) if i >= args.skip]
        if args.only:
            idxs = [i for i in idxs if i in args.only]
        print(f"[-] modo --all: {len(idxs)} bloques -> {idxs}")
    elif args.block is not None:
        if not 0 <= args.block < len(blocks):
            raise SystemExit(f"[!] bloque {args.block} fuera de rango (0-{len(blocks)-1})")
        idxs = [args.block] * max(args.variations, 1)
    else:
        idxs = random.sample(range(len(blocks)), min(args.n, len(blocks)))

    chars = [args.char] * len(idxs) if args.char else \
        [random.choice(usable) for _ in idxs]

    OUT = out_dir(args.pool)
    OUT.mkdir(parents=True, exist_ok=True)
    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    manifest = []
    jobs = []
    for bi, ck in zip(idxs, chars):
        for v in range(args.variations):
            jobs.append((bi, ck, v))

    for i, (bi, ck, v) in enumerate(jobs, 1):
        char = B.get_character(ck)
        if not char:
            print(f"[!] {ck} no esta en personajes.txt, saltando")
            continue
        char_line = char["base"]
        if args.lora_weight is not None:
            char_line = re.sub(r"(<lora:[^:>]+):[^>]*>",
                               rf"\g<1>:{args.lora_weight}>", char_line)
        block = blocks[bi]
        pos = build_prompt(B.master_prompt, char_line, block,
                           cam_brake=not args.no_cam_brake,
                           age_up=args.age_up,
                           force_nude=args.force_nude)
        neg = B.negative_prompt
        tag = args.pool
        tag += "_sincam" if args.no_cam_brake else ""
        name = f"{i:02d}_b{bi:02d}_{ck[:14]}_{tag}_v{v+1}.png"
        src = find_original(block)

        manifest.append({
            "file": name, "block_idx": bi, "char": ck, "pool": args.pool,
            "block": block, "original_txt": str(src) if src else None,
            "prompt": pos,
        })
        print(f"\n[{i}/{len(jobs)}] {name}")
        print(f"    bloque  : {block[:100]}...")
        print(f"    char    : {ck}")
        if args.lora_weight:
            print(f"    lora w  : {args.lora_weight}")
        if src:
            print(f"    original: {src.name}  <- comparar con este")

        if args.dry_run:
            print(f"    prompt  : {pos[:150]}...")
            continue

        payload = {
            "prompt": pos, "negative_prompt": neg,
            "width": 832, "height": 1216, "steps": 28, "cfg_scale": 5.0,
            "sampler_name": "DPM++ 2M Karras", "batch_size": 1, "n_iter": 1,
            "seed": args.seed, "clip_skip": 2,
        }
        try:
            t0 = time.time()
            resp = requests.post(f"{FORGE}/sdapi/v1/txt2img", json=payload, timeout=300)
            resp.raise_for_status()
            data = resp.json()
            (OUT / name).write_bytes(base64.b64decode(data["images"][0]))
            try:
                sd = json.loads(data.get("info", "{}")).get("seed", "?")
            except Exception:
                sd = "?"
            print(f"    [OK] {time.time()-t0:.0f}s seed={sd}")
        except Exception as e:
            print(f"    [FAIL] {e}")

    (OUT / f"manifest_{stamp}.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=1), encoding="utf-8")
    print(f"\n[OK] {len(manifest)} en {OUT}")
    print(f"[OK] manifest_{stamp}.json (bloque + prompt exacto + .txt original)")


if __name__ == "__main__":
    main()