"""
runner: ejecuta un plan del planner y escribe los artefactos del lote.

Mismos artefactos que `batch_poses_chars` (para no tener dos verdades):

    Output/ui_runs/<stamp>/
        <clave>/img/*.png
        <clave>/pack_manifest.json
        <clave>/metrics.json
        manifest_all.json
        plan.json
        ui_spec.json
        audit_1.jpg ... + contact_all.jpg

La generacion NO se reimplementa: es `pack_builder.generate()` job a job. Lo
unico propio es el volcado de artefactos (y las hojas, que aqui van a 6
columnas porque un lote de 45 de un solo personaje no cabe en la hoja de 3).
"""
import json
import random
import shutil
import sys
from datetime import datetime
from pathlib import Path

BASE_DIR = Path(__file__).resolve().parent.parent.parent
for _extra in (BASE_DIR / "Engine", BASE_DIR / "Archivos"):
    if str(_extra) not in sys.path:
        sys.path.insert(0, str(_extra))

import pack_builder as pb  # noqa: E402
import batch_poses_chars as bpc  # noqa: E402
from prompt_builder import PromptBuilder  # noqa: E402
from planner import fs_name  # noqa: E402
from stages import STAGES  # noqa: E402

UI_RUNS = BASE_DIR / "Output" / "ui_runs"
LEVELS = ("sugerente", "spicy", "explicito")

STAGE_ORDER = {stage: i for i, stage in enumerate(STAGES)}


def _safe_dir(path: Path) -> Path:
    """Crea la carpeta; si ya tiene algo dentro, le anade sufijo (no pisa)."""
    if not path.exists():
        path.mkdir(parents=True)
        return path
    if not any(path.iterdir()):
        return path
    n = 2
    while True:
        alt = path.with_name("%s_%d" % (path.name, n))
        if not alt.exists():
            alt.mkdir(parents=True)
            return alt
        n += 1


def _write_json(path: Path, data) -> None:
    path.write_text(json.dumps(data, ensure_ascii=False, indent=1),
                    encoding="utf-8")


def write_sheets(run_dir, records, cols=6, per_sheet=3, thumb=320) -> list:
    """Hojas de contacto: bloques de `per_sheet` personajes, `cols` por fila.

    `audit_N.jpg` por bloque y `contact_all.jpg` con todo. Devuelve las rutas.
    """
    try:
        from PIL import Image, ImageDraw
    except ImportError:
        return []

    run_dir = Path(run_dir)
    by_char = {}
    order = []
    for rec in records:
        char = rec.get("char_dir")
        if char not in by_char:
            by_char[char] = []
            order.append(char)
        by_char[char].append(rec)
    for char in order:
        by_char[char].sort(key=lambda r: STAGE_ORDER.get(r.get("stage"), 99))

    thumb_h = int(thumb * 1.4625)

    def sheet(chars, out_path):
        blocks = [(ch, by_char[ch], 0) for ch in chars if by_char.get(ch)]
        if not blocks:
            return None
        # Columnas utiles: un lote de 3 por personaje no debe salir en 6 columnas
        # (media hoja en blanco). Un pack de 45 si aprovecha las 6.
        widest = max(len(recs) for _, recs, _ in blocks)
        use_cols = max(1, min(cols, widest))
        blocks = [(ch, recs, (len(recs) + use_cols - 1) // use_cols)
                  for ch, recs, _ in blocks]
        width = use_cols * thumb
        height = sum(rows * (thumb_h + 18) for _, _, rows in blocks)
        canvas = Image.new("RGB", (width, height), "white")
        draw = ImageDraw.Draw(canvas)
        y = 0
        for char, recs, rows in blocks:
            for i, rec in enumerate(recs):
                path = run_dir / rec["file"]
                if not path.exists():
                    continue
                try:
                    im = Image.open(path)
                    im.thumbnail((thumb, thumb_h))
                    im = im.convert("RGB")
                except Exception:
                    continue
                x = (i % use_cols) * thumb
                yy = y + (i // use_cols) * (thumb_h + 18)
                canvas.paste(im, (x, yy))
                draw.text((x + 4, yy + thumb_h + 2),
                          "%s | %s %s" % ((rec.get("char") or char)[:22],
                                          rec.get("stage") or
                                          rec.get("nivel"), rec.get("pose_id")),
                          fill="black")
            y += rows * (thumb_h + 18)
        canvas.save(out_path, quality=88)
        return out_path

    out = []
    n = 0
    for i in range(0, len(order), per_sheet):
        n += 1
        got = sheet(order[i:i + per_sheet], run_dir / ("audit_%d.jpg" % n))
        if got:
            out.append(str(got))
    got = sheet(order, run_dir / "contact_all.jpg")
    if got:
        out.append(str(got))
    return out


def run_plan(plan, on_event=None, should_stop=None, resume=False) -> dict:
    """Genera los jobs del plan. -> resumen de la corrida.

    `on_event` recibe dicts con `event` (`char_start`, `image`, `char_done`,
    `done`) y `should_stop()` se consulta entre imagenes para cancelar.
    """
    run_dir = _safe_dir(Path(plan["run_dir"]))
    master = PromptBuilder(str(BASE_DIR))
    items = plan["items"]
    spec = plan["spec"]
    seed = spec.get("seed")
    _write_json(run_dir / "ui_spec.json",
                {**spec, "pick_seed": plan.get("pick_seed"),
                 "quotas": plan.get("quotas"), "run_dir": str(run_dir)})

    total_images = sum(len(item["jobs"]) for item in items)
    all_records, plan_out = [], []
    stopped = False

    for n, item in enumerate(items, 1):
        key = item["char"]
        jobs = item["jobs"]
        folder = item.get("dir") or fs_name(key)
        char_dir = run_dir / folder
        img_dir = char_dir / "img"
        ids = [job.pose_id for job in jobs]

        if on_event:
            on_event({"event": "char_start", "char": key, "dir": folder,
                      "index": n, "total": len(items), "images": len(jobs)})

        records = bpc.load_done(char_dir, ids) if resume else None
        if records is not None:
            metrics = json.loads((char_dir / "metrics.json").read_text(
                encoding="utf-8")) if (char_dir / "metrics.json").exists() else {}
        else:
            budget = max(2 * len(jobs), 4)
            records, metrics = pb.generate(
                jobs, item["char_line"], master.master_prompt,
                master.negative_prompt, img_dir, seed, budget, char=key,
                on_image=(lambda ev: on_event(dict(ev, char=key, dir=folder))
                          if on_event else None),
                should_stop=should_stop)
            char_dir.mkdir(parents=True, exist_ok=True)
            _write_json(char_dir / "pack_manifest.json", records)
            _write_json(char_dir / "metrics.json", metrics)
            if metrics.get("stopped"):
                stopped = True

        for rec in records:
            all_records.append({"char_dir": folder, **rec,
                                "file": "%s/%s" % (folder, rec["file"])})

        poses = {}
        for job in jobs:
            poses.setdefault(job.stage or job.nivel, []).append(
                {"id": job.pose_id, "line": job.pose_line})
        plan_out.append({
            "char": key,
            "char_dir": folder,
            "theme": item["theme"],
            "stages": item["stages"],
            "poses": poses,
            "jobs": [{"nivel": rec["nivel"], "stage": rec.get("stage"),
                      "pose_id": rec["pose_id"], "camera": rec["camera"],
                      "dress": rec["dress"], "seed": rec["seed"],
                      "file": rec["file"]} for rec in records],
        })

        if on_event:
            on_event({"event": "char_done", "char": key, "index": n,
                      "total": len(items), "saved": len(records)})
        if stopped:
            break

    _write_json(run_dir / "manifest_all.json", all_records)
    _write_json(run_dir / "plan.json",
                {"generated": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                 "chars": len(plan_out), "images": total_images,
                 "seed": seed, "pick_seed": plan.get("pick_seed"),
                 "stages": spec.get("stages"), "mode": spec.get("mode"),
                 "plan": plan_out})

    sheets = []
    if all_records:
        try:
            sheets = write_sheets(run_dir, all_records)
        except Exception as exc:      # las hojas no pueden tumbar un lote
            print("[!] hojas de contacto: %s" % exc)

    summary = {
        "run_dir": str(run_dir),
        "planned": total_images,
        "generated": len(all_records),
        "chars": len(plan_out),
        "stopped": stopped,
        "sheets": sheets,
    }
    _write_json(run_dir / "summary.json", summary)
    if on_event:
        on_event({"event": "done", **summary})
    return summary


# ------------------------------------------------------------- regenerar 1

def find_record(all_records, char, pose_id=None, stage=None, file=None):
    for rec in all_records:
        # `char` es la clave (identidad) y `char_dir` la carpeta; hay lotes
        # viejos donde solo existe `char_dir`, asi que se aceptan las dos.
        if rec.get("char_dir") != char and rec.get("char") != char:
            continue
        if file and rec.get("file") == file:
            return rec
        if pose_id and rec.get("pose_id") == pose_id:
            if stage and rec.get("stage") != stage:
                continue
            return rec
    return None


def regen_one(run_dir, char, pose_id=None, stage=None, file=None, seed=None,
              budget=4, neg_extra=None, dry_run=False) -> dict:
    """Regenera UNA imagen: mismo job, misma camara, semilla nueva.

    Reconstruye el job desde `plan.json`/`manifest_all.json` (el prompt sale
    de la linea ACTUAL del banco y del outfit ACTUAL del tema) y compara el
    prompt nuevo con el guardado: si cambia, se avisa (`drift`) pero manda el
    nuevo. Lo viejo se mueve a `<char>/_reemplazadas/`.
    """
    run_dir = Path(run_dir)
    plan_path = run_dir / "plan.json"
    all_path = run_dir / "manifest_all.json"
    if not plan_path.exists() or not all_path.exists():
        raise ValueError("%s no parece un lote (falta plan.json o "
                         "manifest_all.json)" % run_dir)
    plan_doc = json.loads(plan_path.read_text(encoding="utf-8"))
    all_records = json.loads(all_path.read_text(encoding="utf-8"))

    rec = find_record(all_records, char, pose_id, stage, file)
    if not rec:
        raise ValueError("no encuentro esa imagen en el lote "
                         "(char=%s pose=%s stage=%s)" % (char, pose_id, stage))
    folder = rec.get("char_dir") or fs_name(char)

    B = PromptBuilder(str(BASE_DIR))
    import pose_bank
    bank = pose_bank.load_bank()
    avoid = pose_bank.load_avoid()
    cameras = pb.load_cameras()
    pools = {lv: pb.load_pool_lines(lv) for lv in LEVELS}
    neg_base = B.negative_prompt
    if neg_extra:
        neg_base = "%s, %s" % (neg_base, neg_extra) if neg_base else neg_extra

    char_data = B.get_character(char)
    if not char_data:
        raise ValueError("'%s' no esta en personajes.txt" % char)

    nivel = rec["nivel"]
    quotas = {lv: (1 if lv == nivel else 0) for lv in LEVELS}
    variants = {k: v for k, v in (rec.get("variants") or {}).items() if v}
    theme = rec.get("theme")
    extra = {}
    if rec.get("stage"):
        extra["stage_quotas"] = {rec["stage"]: 1}
        extra["allow_legacy_outfit"] = True

    jobs = pb.compose_theme_jobs(
        [theme], 1, quotas, "fixed", variants, rec["dress"], 0.0,
        bank, avoid, pools, cameras, random.Random(0),
        max_variants=1, bank_ids=[rec["pose_id"]], **extra)
    if len(jobs) != 1:
        raise ValueError("%s/%s: salieron %d jobs"
                         % (char, rec["pose_id"], len(jobs)))
    job = jobs[0]
    job.camera = rec["camera"]      # con 1 job el indice cambiaria el angulo
    job.seed = seed

    prompt = pb.build_run_prompt(B.master_prompt, char_data["base"], job)
    drift = prompt.strip() != (rec.get("prompt") or "").strip()

    result = {
        "char": char, "nivel": nivel, "stage": rec.get("stage"),
        "pose_id": rec["pose_id"], "old_file": rec["file"],
        "theme": theme, "camera": job.camera, "dress": job.dress,
        "prompt": prompt, "drift": drift, "dry_run": bool(dry_run),
        "new_file": None, "seed": None,
    }
    if dry_run:
        return result

    char_dir = run_dir / folder
    img_dir = char_dir / "img"
    tmp_dir = char_dir / "_regen_tmp"
    if tmp_dir.exists():
        shutil.rmtree(tmp_dir)

    records, _ = pb.generate([job], char_data["base"], B.master_prompt,
                             neg_base, tmp_dir, seed, budget, char=char)
    if not records:
        shutil.rmtree(tmp_dir, ignore_errors=True)
        raise ValueError("no se genero nada para %s/%s"
                         % (char, rec["pose_id"]))

    new_rec = records[0]
    old_name = Path(rec["file"]).name
    new_name = Path(new_rec["file"]).name
    if "_" in old_name and "_" in new_name:
        new_name = "%s_%s" % (old_name.split("_")[0],
                              new_name.split("_", 1)[1])
    shutil.move(str(tmp_dir / Path(new_rec["file"]).name),
                str(img_dir / new_name))
    shutil.rmtree(tmp_dir, ignore_errors=True)

    old_png = img_dir / old_name
    if old_png.exists():
        hist = char_dir / "_reemplazadas"
        hist.mkdir(parents=True, exist_ok=True)
        shutil.move(str(old_png), str(hist / old_name))

    new_rec = dict(new_rec, file="img/%s" % new_name)
    char_man = char_dir / "pack_manifest.json"
    char_recs = json.loads(char_man.read_text(encoding="utf-8"))
    char_recs = [new_rec if (r.get("pose_id") == rec["pose_id"]
                             and r.get("stage") == rec.get("stage"))
                 else r for r in char_recs]
    _write_json(char_man, char_recs)

    all_records = [({**new_rec, "char_dir": folder,
                     "file": "%s/%s" % (folder, new_rec["file"])}
                    if (r.get("char_dir") == folder
                        and r.get("pose_id") == rec["pose_id"]
                        and r.get("stage") == rec.get("stage")) else r)
                   for r in all_records]
    _write_json(all_path, all_records)

    for item in plan_doc.get("plan", []):
        if item.get("char") != char:
            continue
        for entry in item.get("jobs", []):
            if (entry.get("pose_id") == rec["pose_id"]
                    and entry.get("stage") == rec.get("stage")):
                entry.update({"nivel": new_rec["nivel"],
                              "camera": new_rec["camera"],
                              "dress": new_rec["dress"],
                              "seed": new_rec["seed"],
                              "file": new_rec["file"]})
    _write_json(plan_path, plan_doc)

    result["new_file"] = "%s/%s" % (folder, new_rec["file"])
    result["seed"] = new_rec["seed"]
    return result