"""
Regenera los APROBADOS con el ensamblador actual (BREAK + hombre simplificado
+ dedupe de conteos + hombre reducido en POV).

No manipula el prompt viejo con regex: reconstruye desde la linea del personaje
y el bloque, que es lo que hace test_blocks.py. Asi el regenerado y una
generacion normal dan EXACTAMENTE el mismo prompt.

    python Engine/regen_aprobados.py --dry-run
    python Engine/regen_aprobados.py --pool explicit

Salida en Output/test_blocks/<pool>/REGEN/ (mismos nombres que APROBADO/).
"""
import argparse
import base64
import json
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
from test_blocks import assemble_prompt  # noqa: E402

TEST_BLOCKS = BASE_DIR / "Output" / "test_blocks"
FORGE = "http://127.0.0.1:7860"


def load_batch(pool):
    """-> {nombre_png: (char, block)}.

    Fusiona TODOS los manifests: algunas imagenes se generaron en corridas
    viejas cuyo manifest no es el mas grande, asi que si solo se mira el mayor
    se pierden.
    """
    out = {}
    for m in sorted(TEST_BLOCKS.rglob("manifest_*.json")):
        try:
            data = json.loads(m.read_text(encoding="utf-8-sig"))
        except Exception:
            continue
        for e in data:
            f = e.get("file", "")
            if f and f not in out and e.get("char") and e.get("block"):
                out[f] = (e["char"], e["block"])
    return out


def load_current_blocks(pool):
    """-> {.txt de origen: bloque CURADO del pool}.

    Importa: el manifest guarda el bloque TAL COMO ESTABA cuando se genero la
    imagen. Si despues se limpio el bloque (por ejemplo quitando 'solo_focus',
    que rompia la escena), el manifest se queda obsoleto y regenerar desde ahi
    reintroduce el tag que acabamos de eliminar. El pool es la version buena.
    """
    h = BASE_DIR / "Archivos" / "wildcards_harvest"
    idx = h / "bloques_indice.txt"
    if not idx.exists():
        return {}
    per_pool = {}
    for line in idx.read_text(encoding="utf-8").splitlines():
        if "\t" not in line or line.startswith("#"):
            continue
        gi, sev, org = line.split("\t")
        per_pool.setdefault(sev, []).append((int(gi), org.split("/")[-1]))
    out = {}
    for sev, items in per_pool.items():
        f = h / ("bloques_%s.txt" % sev)
        if not f.exists():
            continue
        lines = [l.strip() for l in f.read_text(encoding="utf-8").splitlines()
                 if l.strip() and not l.startswith("#")]
        # OJO: 'gi' es el indice GLOBAL, pero 'lines' es el pool de UNA
        # severidad. Hay que enumerar dentro del pool, no indexar por gi.
        for pos, (gi, name) in enumerate(items):
            if pos < len(lines):
                out[name] = lines[pos]
    return out


def main():
    p = argparse.ArgumentParser(description="Regenera APROBADOS (ensamblador actual)")
    p.add_argument("--pool", default="explicit")
    p.add_argument("--variations", type=int, default=1,
                   help="seeds por imagen")
    p.add_argument("--dry-run", action="store_true")
    p.add_argument("--steps", type=int, default=28)
    p.add_argument("--only", action="append", default=None,
                   help="regenera solo estos nombres de png (repetible)")
    args = p.parse_args()

    appr = TEST_BLOCKS / args.pool / "APROBADO"
    batch = load_batch(args.pool)
    current = load_current_blocks(args.pool)
    B = PromptBuilder(str(BASE_DIR))

    pngs = sorted(f.name for f in appr.glob("*.png"))
    if args.only:
        pngs = [p_ for p_ in pngs if p_ in args.only]

    # los que no tienen char/block en el manifest (corridas viejas)
    origin_of = {}
    for m in sorted(TEST_BLOCKS.rglob("manifest_*.json")):
        try:
            data = json.loads(m.read_text(encoding="utf-8-sig"))
        except Exception:
            continue
        for e in data:
            f = e.get("file", "")
            if f and f not in origin_of:
                origin_of[f] = Path(e.get("original_txt") or "").name

    jobs, skip = [], []
    for png in pngs:
        if png not in batch:
            skip.append(png)
            continue
        ckey, block = batch[png]
        ch = B.get_character(ckey)
        # el bloque del POOL manda: el manifest guarda el bloque tal como estaba
        # al generar, y si despues se limpio (p.ej. quitando 'solo_focus'),
        # regenerar desde el manifest reintroduce el tag eliminado.
        orig = origin_of.get(png)
        if orig and orig in current:
            block = current[orig]
        if not ch or not block:
            skip.append(png)
            continue
        for v in range(max(args.variations, 1)):
            pos, segs = assemble_prompt(B.master_prompt, ch["base"], block)
            jobs.append((png, ckey, block, pos, segs, v))

    print(f"{len(jobs)} a regenerar en {args.pool}/REGEN/")
    if skip:
        print(f"{len(skip)} sin datos en el manifest (no se pueden reconstruir):")
        for s in skip:
            print(f"    {s}")

    if args.dry_run:
        for png, ckey, block, pos, segs, v in jobs:
            print("\n" + "=" * 74)
            print(png + ("  (variacion %d)" % (v + 1) if args.variations > 1 else ""))
            print(pos)
        return

    outdir = TEST_BLOCKS / args.pool / "REGEN"
    outdir.mkdir(parents=True, exist_ok=True)

    import requests
    manifest = []
    for i, (png, ckey, block, pos, segs, v) in enumerate(jobs, 1):
        payload = {"prompt": pos, "negative_prompt": B.negative_prompt,
                   "width": 832, "height": 1216, "steps": args.steps,
                   "cfg_scale": 5.0, "sampler_name": "DPM++ 2M Karras",
                   "batch_size": 1, "n_iter": 1, "seed": -1, "clip_skip": 2}
        # con variaciones el nombre lleva v2, v3... para no pisarse
        stem = png[:-4]
        name = png if args.variations <= 1 else f"{stem}_v{v + 1}.png"
        print(f"[{i}/{len(jobs)}] {name}")
        try:
            t0 = time.time()
            r = requests.post(f"{FORGE}/sdapi/v1/txt2img", json=payload, timeout=300)
            r.raise_for_status()
            d = r.json()
            (outdir / name).write_bytes(base64.b64decode(d["images"][0]))
            try:
                seed = json.loads(d.get("info", "{}")).get("seed", "?")
            except Exception:
                seed = "?"
            print(f"    [OK] {time.time()-t0:.0f}s seed={seed} BREAKs={pos.count('BREAK')}")
            manifest.append({"file": name, "char": ckey, "block": block,
                             "seed": seed, "prompt": pos, "segments": segs,
                             "variacion": v + 1})
        except Exception as e:
            print(f"    [FAIL] {e}")

    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    (outdir / f"regen_{stamp}.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=1), encoding="utf-8")
    print(f"\n[OK] {len(manifest)} -> {outdir}")
    print(f"[OK] regen_{stamp}.json (prompt + segmentos + block)")


if __name__ == "__main__":
    main()