"""
pixiv_per_image (D1.1): los 2-3 tags de Pixiv que cambian segun la imagen.

`pixiv_tags.py` da la cadena del personaje (`fijos` + `por_imagen`), pero
`por_imagen` sale de lo que mas co-ocurre con el personaje: son los mismos tags
para sus 3 imagenes del lote. Esto los sustituye leyendo el prompt real de cada
PNG (`sex_from_behind`, `pantyhose`, `breasts_out`...) y mapeandolo al tag
japones que existe de verdad en Pixiv (`Archivos/pixiv_tag_map.json`).

    python Engine/pixiv_per_image.py                      # ultimo BATCH_* con manifest_all.json
    python Engine/pixiv_per_image.py --run "<run>"        # un lote concreto
    python Engine/pixiv_per_image.py --n 2                # 2 tags por imagen (default 3)
    python Engine/pixiv_per_image.py --char "Furina,Yelan"
    python Engine/pixiv_per_image.py --dry-run            # imprime, no escribe
    python Engine/pixiv_per_image.py --r18-all            # R-18 tambien en las SFW
    python Engine/pixiv_per_image.py --add "03_explicito=精液"   # lo que solo se
                                                          # ve en la auditoria

Escribe en el lote:
    pixiv_tags.json   una entrada por imagen, con `tags`, `extra` y el origen de cada tag
    pixiv_tags.md     las cadenas listas para copiar/pegar, agrupadas por personaje

Reglas (D1.3): 10 tags = personaje JP+EN + obra JP+EN + 5 NSFW/escena + 1 libre.
Los 5 NSFW/escena son 2-3 de cuerpo (fijos del personaje) + `--n` de la imagen.
`R-18` y `AIイラスト` van fuera de la cuenta. Una imagen vestida y sin acto ni
estado explicito NO lleva `R-18` ni tags de acto: se marca como SFW.
"""
import argparse
import json
import re
import sys
from datetime import datetime
from pathlib import Path

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

BASE_DIR = Path(__file__).resolve().parent.parent
POOL = BASE_DIR / "Archivos" / "pixiv_nsfw_pool.json"
TMAP = BASE_DIR / "Archivos" / "pixiv_tag_map.json"
PILOTS = BASE_DIR / "Output" / "pilots"

SLOTS = 10          # limite de Pixiv
FREE = 1            # el "libre" que se reserva al final
EXTRA = ("R-18", "AIイラスト")

# De que depende mas el tag: primero lo que hace la imagen (acto), luego como
# esta (desnuda/semidesnuda/semen), luego la ropa, el encuadre y por ultimo el
# cuerpo (que casi siempre ya viene en los fijos del personaje).
PRIO = {"acto": 0, "estado": 1, "ropa": 2, "encuadre": 3, "cuerpo": 4}

# Dentro de la ropa, lo especifico manda sobre lo generico: si el prompt dice
# `swimsuit` y `panties`, el tag bueno es 水着, no 下着 (y no caben los dos).
ROPA_ORDER = ["水着", "ビキニ", "マイクロビキニ", "競泳水着", "レオタード",
              "制服", "体操服", "チャイナドレス", "裸エプロン", "ボンデージ",
              "黒タイツ", "ニーソ", "網タイツ", "ガーター", "透け", "下着"]
ROPA_RANK = {t: i for i, t in enumerate(ROPA_ORDER)}

# Si la imagen no da 3 tags utiles y no dice nada de ropa, se dice en que estado
# esta. `着衣` NO es NSFW: una imagen vestida y sin acto se publica como SFW.
DRESS = {"dressed": "着衣", "stripping": "半裸", "nude": "全裸"}
NSFW_STATE = {"全裸", "半裸", "精液", "溢れ精液", "アヘ顔", "絶頂", "愛液",
              "ボテ腹", "妊娠", "失神"}

# Tags que no pueden ir juntos sin contradecirse: no se ponen dos "versiones"
# de la misma prenda. El primero que aparece manda.
SWIM = {"水着", "ビキニ", "マイクロビキニ", "競泳水着", "レオタード"}


def conflicts(tag: str, chosen) -> bool:
    if tag == "下着":
        return any(c in SWIM for c in chosen)
    if tag in SWIM:
        return any(c == "下着" or (c in SWIM and c != tag) for c in chosen)
    if tag == "着衣":
        return any(c in ("全裸", "半裸") for c in chosen)
    if tag in ("全裸", "半裸"):
        return "着衣" in chosen
    return False

# Ruido del ensamblador: no describe la imagen.
NOISE = re.compile(
    r"^(masterpiece|best quality|highly detailed|sharp focus|very aesthetic|"
    r"very awa|detailed background|score_\d+(_up)?|absurdres|highres|"
    r"blurry background|simple background)$", re.I)


def load_json(path: Path, default=None):
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except FileNotFoundError:
        return default if default is not None else {}
    except json.JSONDecodeError as e:
        raise SystemExit(f"[!] {path} no es JSON valido: {e}")


def norm_tag(raw: str) -> str:
    """Tag del prompt -> forma comparable: minusculas, sin pesos ni espacios."""
    t = (raw or "").strip().lower()
    t = re.sub(r"<lora:[^>]*>", "", t)
    t = re.sub(r"\(([^():]+):[0-9.]+\)", r"\1", t)      # (faceless male:1.4)
    t = re.sub(r"[\[\]{}]", "", t)
    t = t.replace("\\)", ")").replace("\\(", "(")
    t = re.sub(r"\s+", "_", t)
    t = re.sub(r"_+", "_", t).strip("_ ,")
    return t


def tokens(prompt: str) -> list:
    """Tags del prompt en orden, sin BREAK ni pesos (los repetidos se dejan)."""
    out = []
    for block in (prompt or "").split("BREAK"):
        for chunk in block.split(","):
            t = norm_tag(chunk)
            if not t or NOISE.match(t.replace("_", " ")):
                continue
            out.append(t)
    return out


def find_run(explicit=None) -> Path:
    """Carpeta del lote: la que se pase o el BATCH_* mas reciente con manifest."""
    if explicit:
        p = Path(explicit)
        if p.is_file():
            p = p.parent
        if not (p / "manifest_all.json").exists():
            raise SystemExit(f"[!] {p} no tiene manifest_all.json")
        return p.resolve()
    found = sorted(PILOTS.rglob("BATCH_*/manifest_all.json"),
                   key=lambda f: f.stat().st_mtime)
    if not found:
        raise SystemExit(f"[!] no hay ningun Output/pilots/**/BATCH_*/manifest_all.json")
    return found[-1].parent.resolve()


def chain_for(key: str) -> dict | None:
    """Ficha de `Archivos/pixiv_tags/<clave>.json` (misma ruta que pixiv_tags)."""
    from pixiv_tags import chain_path      # una sola definicion de la ruta
    path = chain_path(key)
    if not path.exists():
        return None
    doc = load_json(path, {})
    if not doc.get("fijos"):
        return None
    return doc


def split_fijos(fijos: list, pool: dict) -> tuple:
    """(identidad, cuerpo): los de cuerpo son los que se recortan si falta sitio."""
    body = [t for t in fijos if (pool.get(t) or {}).get("cat") == "cuerpo"]
    head = [t for t in fijos if t not in body]
    return head, body


def mapped(rec: dict, fijos: list, tmap: dict, pool: dict) -> list:
    """Candidatos del prompt -> tag de Pixiv, ordenados por PRIO. Sin duplicados."""
    seen, out = set(), []
    for tok in tokens(rec.get("prompt", "")):
        jp = tmap.get(tok)
        if not jp or jp in seen or jp in fijos:
            continue
        seen.add(jp)
        out.append({"pixiv": jp, "from": tok,
                    "cat": (pool.get(jp) or {}).get("cat", "")})
    return sorted(out, key=lambda d: (PRIO.get(d["cat"], 9),
                                      ROPA_RANK.get(d["pixiv"], 99)))


def build_image(rec: dict, chain: dict, tmap: dict, pool: dict, n: int,
                r18_all: bool = False, manual: list | None = None) -> dict:
    """La cadena de UNA imagen: identidad + tags de la imagen + cuerpo + libre."""
    fijos = list(chain.get("fijos") or [])
    head, body = split_fijos(fijos, pool)
    used = set(fijos)

    added = []
    for cand in mapped(rec, used, tmap, pool):
        if len(added) >= n:
            break
        if cand["pixiv"] in used or conflicts(cand["pixiv"], used):
            continue
        added.append(cand)
        used.add(cand["pixiv"])

    # Correcciones de la auditoria visual (`--add`): lo que se ve en el PNG y el
    # prompt no puede saber, tipicamente el semen que el checkpoint anade solo.
    for jp in (manual or []):
        if len(added) >= 5 or jp in used or conflicts(jp, used):
            continue
        added.append({"pixiv": jp, "from": "manual",
                      "cat": (pool.get(jp) or {}).get("cat", "")})
        used.add(jp)

    # Ropa/estado de la imagen si el prompt no daba para 3: mejor decir "esta
    # vestida" que inventar un acto que no aparece en el dibujo. Pero si ya hay
    # ropa (制服, 下着, ビキニ) o estado (半裸), no se anade: se contradirian.
    dress = (rec.get("dress") or "").strip().lower()
    dress_tag = DRESS.get(dress)
    hay_ropa = any(a["cat"] in ("ropa", "estado") for a in added)
    if len(added) < n and dress_tag and not hay_ropa and dress_tag not in used \
            and not conflicts(dress_tag, used):
        added.append({"pixiv": dress_tag, "from": f"dress={dress}",
                      "cat": (pool.get(dress_tag) or {}).get("cat", "estado")})
        used.add(dress_tag)

    per = [a["pixiv"] for a in added]
    keep = max(0, min(len(body), SLOTS - FREE - len(head) - len(per)))
    tags = (head + per + body[:keep] + [""] * SLOTS)[:SLOTS]

    cats = {a["cat"] for a in added}
    explicit_state = any(a["cat"] == "estado" and a["pixiv"] in NSFW_STATE
                         for a in added)
    nsfw = dress in ("stripping", "nude") or "acto" in cats or explicit_state
    nsfw = nsfw or r18_all
    extra = [t for t in EXTRA if nsfw or t != "R-18"]
    line = ", ".join([t for t in tags if t] + extra)

    warnings = []
    if len([t for t in tags if t]) < SLOTS - 1:
        warnings.append(f"solo {len([t for t in tags if t])} tags utiles")
    if not nsfw:
        warnings.append("SFW: sin R-18 ni tags de acto (el prompt no los da)")

    return {
        "file": rec.get("file"),
        "char": rec.get("char"),
        "nivel": rec.get("nivel"),
        "tema": rec.get("tema"),
        "seed": rec.get("seed"),
        "dress": dress or None,
        "tags": tags,
        "extra": extra,
        "line": line,
        "added": added,
        "base_body": body[:keep],
        "nsfw": nsfw,
        "warnings": warnings,
    }


def build_markdown(doc: dict) -> str:
    out = [f"# Tags de Pixiv — `{doc['run']}`", "",
           f"Generado {doc['generated']} · **{doc['n']}** tags por imagen · "
           f"{len(doc['images'])} imagenes.", ""]
    if doc["missing"]:
        out += ["> Sin ficha de personaje (no hay cadena): "
                + ", ".join(f"`{k}`" for k in doc["missing"]), ""]

    by_char = {}
    for img in doc["images"]:
        by_char.setdefault(img["char"], []).append(img)

    for key, imgs in by_char.items():
        info = doc["chars"].get(key) or {}
        out.append(f"## {key} — {info.get('char_jp', '?')} / {info.get('char_en', '?')}")
        out.append(f"Obra: {info.get('series_jp', '?')} / {info.get('series_en', '?')}"
                   f" · ficha: `Archivos/pixiv_tags/{key}.json`")
        out.append("")
        for img in imgs:
            name = Path(img["file"] or "?").name
            tema = ", ".join(img["tema"] or []) if isinstance(img["tema"], list) else img["tema"]
            out.append(f"### {name}")
            out.append(f"**{img['nivel']}** · tema `{tema}` · seed `{img['seed']}`"
                       f" · ropa **{img['dress'] or '?'}**"
                       f" · {'NSFW' if img['nsfw'] else 'SFW'}")
            out.append("")
            out.append("```")
            out.append(img["line"])
            out.append("```")
            if img["added"]:
                out.append("")
                for a in img["added"]:
                    out.append(f"- `{a['pixiv']}` ← `{a['from']}` ({a['cat'] or '?'})")
            if img["base_body"]:
                out.append("")
                out.append("Cuerpo del personaje: "
                           + ", ".join(f"`{t}`" for t in img["base_body"]))
            for w in img["warnings"]:
                out.append("")
                out.append(f"> {w}")
            out.append("")
    return "\n".join(out)


def main() -> None:
    p = argparse.ArgumentParser(
        description="Tags de Pixiv por imagen (D1.1) leyendo el prompt del lote")
    p.add_argument("--run", help="carpeta del lote (default: el BATCH_* mas reciente)")
    p.add_argument("--n", type=int, default=3,
                   help="tags que se anaden por imagen (default 3)")
    p.add_argument("--char", help="solo estos personajes (separados por comas)")
    p.add_argument("--dry-run", action="store_true",
                   help="imprime el markdown, no escribe nada en el lote")
    p.add_argument("--r18-all", action="store_true",
                   help="pone R-18 tambien en las imagenes SFW (por defecto solo "
                        "en las que tienen acto/estado explicito)")
    p.add_argument("--add", action="append", default=[], metavar="FILTRO=TAG[,TAG]",
                   help="tags a mano para las imagenes cuyo nombre contenga FILTRO "
                        "(repetible). Sirve para lo que se ve en el PNG pero no "
                        "esta en el prompt, p. ej. 精液")
    args = p.parse_args()

    pool = load_json(POOL, {})
    tmap = load_json(TMAP, {})
    if not pool:
        raise SystemExit("[!] falta Archivos/pixiv_nsfw_pool.json "
                         "(genera: python Engine/pixiv_nsfw_pool.py)")
    if not tmap:
        raise SystemExit("[!] falta Archivos/pixiv_tag_map.json")
    if args.n < 0 or args.n > 5:
        raise SystemExit("[!] --n razonable: 0..5 (10 tags los comparten todos)")

    run = find_run(args.run)
    records = load_json(run / "manifest_all.json", [])
    if not isinstance(records, list) or not records:
        raise SystemExit(f"[!] {run / 'manifest_all.json'} no tiene imagenes")

    wanted = None
    if args.char:
        wanted = {c.strip() for c in args.char.split(",") if c.strip()}

    manual = []
    for spec in args.add:
        if "=" not in spec:
            raise SystemExit(f"[!] --add espera FILTRO=TAG[,TAG], no {spec!r}")
        needle, tags = spec.split("=", 1)
        tags = [t.strip() for t in tags.split(",") if t.strip()]
        if not needle.strip() or not tags:
            raise SystemExit(f"[!] --add incompleto: {spec!r}")
        manual.append((needle.strip(), tags))

    images, missing, chars = [], [], {}
    for rec in records:
        key = rec.get("char")
        if not key or (wanted and key not in wanted):
            continue
        chain = chain_for(key)
        if chain is None:
            if key not in missing:
                missing.append(key)
            continue
        img = build_image(rec, chain, tmap, pool, args.n, args.r18_all,
                          [t for needle, tags in manual
                           if needle.lower() in (rec.get("file") or "").lower()
                           for t in tags])
        images.append(img)
        chars.setdefault(key, {
            "char_jp": (chain.get("jp") or {}).get("char", ""),
            "char_en": (chain.get("en") or {}).get("char", ""),
            "series_jp": (chain.get("jp") or {}).get("series", ""),
            "series_en": (chain.get("en") or {}).get("series", ""),
        })

    doc = {
        "run": str(run),
        "generated": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "n": args.n,
        "images": images,
        "chars": chars,
        "missing": missing,
    }
    md = build_markdown(doc)

    if args.dry_run:
        print(md)
        print(f"[--] dry-run: {len(images)} imagenes, {len(missing)} sin ficha. "
              f"No se ha escrito nada.")
        return

    js = run / "pixiv_tags.json"
    mdp = run / "pixiv_tags.md"
    js.write_text(json.dumps(doc, ensure_ascii=False, indent=1), encoding="utf-8")
    mdp.write_text(md, encoding="utf-8")

    nsfw = sum(1 for i in images if i["nsfw"])
    print(f"[OK] {len(images)} imagenes ({nsfw} NSFW / {len(images) - nsfw} SFW) "
          f"de {len(chars)} personajes")
    for i in images:
        extra = ", ".join(i["added"] and [a["pixiv"] for a in i["added"]] or []) or "—"
        print(f"  {Path(i['file']).name:<44} {i['nivel']:<9} + {extra}")
    if missing:
        print(f"[!] sin ficha en Archivos/pixiv_tags (no se inventa): {', '.join(missing)}")
    print(f"[OK] {js}")
    print(f"[OK] {mdp}")


if __name__ == "__main__":
    main()