"""
Descarga referencias de poses de Danbooru CON sus tags, para validar el
pipeline de bloques contra la realidad.

Por que esto sirve: Danbooru ya tiene las etiquetas hechas por la comunidad y
el campo 'rating' (g/s/q/e) resuelto. Eso contrasta con dos cosas que aqui son
suposiciones nuestras:
  - el 'rating' de origen frente a nuestra clasificacion de severidad por tags
  - los tags de Danbooru frente a lo que devuelve WD14 sobre la misma imagen

Descarga en Archivos/wd14/danbooru/ (NO en inbox/, que ya esta tagueada).
Cada imagen viene con un .txt al lado y todo queda en un manifiesto.

    python Engine/danbooru_refs.py --limit 7
    python Engine/danbooru_refs.py --limit 7 --dry-run

Despues:
    python Engine/wd14_tagger.py --folder Archivos/wd14/danbooru \\
        --out Archivos/wd14/tags_danbooru
    python Engine/comparar_danbooru.py       <- contrasta ambas fuentes
"""
import argparse
import json
import time
from datetime import datetime
from pathlib import Path

import requests

BASE_DIR = Path(__file__).resolve().parent.parent
DEST = BASE_DIR / "Archivos" / "wd14" / "danbooru"
API = "https://danbooru.donmai.us/posts.json"
UA = "plumloom-research/1.0 (personal pose research; local)"

# (tag de pose, rating). Se reparten para cubrir las tres severidades.
# Los de rating 'g' son poses seguras; 'e' son las que rompen anatomia.
BUSQUEDAS = [
    ("standing", "g"), ("sitting", "g"), ("lying", "g"), ("kneeling", "g"),
    ("from_behind", "s"), ("bent_over", "s"), ("spread_legs", "s"),
    ("looking_back", "q"), ("all_fours", "q"), ("grabbing", "q"),
    ("doggystyle", "e"), ("prone_bone", "e"), ("cowgirl_position", "e"),
    ("reverse_cowgirl_position", "e"), ("missionary", "e"),
]

# Solo imagenes. Danbooru devuelve tambien .mp4/.webm/.gif y un solo clip
# puede pesar 30 MB (pasó: 28,4 MB de un reverse_cowgirl_position).
IMG_EXT = (".jpg", ".jpeg", ".png", ".webp", ".bmp")

# Tags que delatan una captura de juego, un comic o una repost con marca:
# no sirven como referencia de pose.
BASURA = ("gameplay_mechanics", "multiple_views", "comic", "panel",
          "commentary_request", "translated", "untranslatable_commentary",
          "english_text", "chinese_text", "korean_text", "scan",
          "advertisement", "product_placement", "character_sheet",
          "height_chart", "composite", "loop", "animated",
          "watermark", "weibo_username", "weibo_watermark", "pixiv_id",
          "signature", "username", "artist_request", "commission")

POSTS = ("post", "posts")
POST_FIELDS = ("id", "md5", "rating", "tag_count", "score", "fav_count",
               "is_pending", "is_flagged", "is_banned", "file_ext",
               "image_width", "image_height", "large_file_url", "file_url",
               "tag_string_general", "tag_string_character",
               "tag_string_copyright", "tag_string_artist", "source")


def buscar(ses, tag, rating, n):
    """n posts de Danbooru para ese tag y rating, ya filtrados."""
    url = "%s?tags=%s%%20rating:%s&limit=%d" % (API, tag, rating, n * 3)
    r = ses.get(url, timeout=30)
    r.raise_for_status()
    out = []
    for p in r.json():
        if p.get("is_flagged") or p.get("is_banned"):
            continue
        gen = (p.get("tag_string_general") or "").split()
        if len(gen) > 45:                       # capturas y compositions
            continue
        if sum(1 for t in gen if t in BASURA) >= 3:
            continue
        if p.get("score", 0) < 2 and p.get("fav_count", 0) < 2:
            continue
        out.append(p)
        if len(out) >= n:
            break
    return out


def main():
    ap = argparse.ArgumentParser(description="Descarga referencias de Danbooru")
    ap.add_argument("--limit", type=int, default=7,
                    help="imagenes por busqueda (x15 búsquedas)")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    ses = requests.Session()
    ses.headers["User-Agent"] = UA

    if not args.dry_run:
        DEST.mkdir(parents=True, exist_ok=True)

    visto, filas = set(), []
    for tag, rating in BUSQUEDAS:
        try:
            posts = buscar(ses, tag, rating, args.limit)
        except Exception as e:
            print("  [FALLO] %s rating:%s -> %s" % (tag, rating, e))
            continue
        n_ok = 0
        for p in posts:
            if p["id"] in visto:
                continue
            if p.get("is_flagged") or p.get("is_banned"):
                continue
            # fuera lo que no es imagen (mp4/webm/gif): pesa 30 MB y WD14 no lo lee
            if ("." + (p.get("file_ext") or "").lower()) not in IMG_EXT:
                continue
            visto.add(p["id"])
            n_ok += 1
            filas.append((tag, rating, p))
        print("  %-26s rating:%s  %d utiles" % (tag, rating, n_ok))

    print("\n  total unicos: %d" % len(filas))
    if args.dry_run:
        for tag, rating, p in filas[:12]:
            gen = (p["tag_string_general"] or "").split()
            print("\n  [%s rating:%s] %d tags, %dx%d, autor=%s" % (
                tag, rating, p["tag_count"], p["image_width"],
                p["image_height"], p["tag_string_artist"] or "?"))
            print("    " + " ".join(gen[:26]))
        return

    manifest = []
    for i, (tag, rating, p) in enumerate(filas, 1):
        # la version 'sample' (850px) basta de sobra para taguear y pesa poco
        url = p.get("large_file_url") or p.get("file_url")
        ext = p.get("file_ext") or "jpg"
        name = "%s_%s_%s.%s" % (p["id"], tag, rating, ext)
        dest = DEST / name
        try:
            r = ses.get(url, timeout=60)
            r.raise_for_status()
            dest.write_bytes(r.content)
        except Exception as e:
            print("  [FALLO descarga] %s -> %s" % (name, e))
            continue

        # los tags de Danbooru al lado, ordenados como los deja la comunidad
        txt = dest.with_suffix(".txt")
        txt.write_text(" ".join((p["tag_string_general"] or "").split()),
                       encoding="utf-8")

        manifest.append({
            "archivo": name, "post_id": p["id"], "busqueda": tag,
            "rating": rating, "score": p["score"], "fav_count": p["fav_count"],
            "tag_count": p["tag_count"],
            "resolucion": "%dx%d" % (p["image_width"], p["image_height"]),
            "artist": p["tag_string_artist"], "character": p["tag_string_character"],
            "copyright": p["tag_string_copyright"], "source": p.get("source"),
            "url": url,
            "danbooru_tags": (p["tag_string_general"] or "").split(),
        })
        print("  [%2d/%2d] %-34s %6.0f KB" % (
            i, len(filas), name, len(dest.read_bytes()) / 1024))
        time.sleep(0.4)          # no saturar la API

    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    man = DEST / f"danbooru_manifest_{stamp}.json"
    man.write_text(json.dumps(manifest, ensure_ascii=False, indent=1),
                   encoding="utf-8")

    print("\n[OK] %d imagenes -> %s" % (len(manifest), DEST))
    print("[OK] %s" % man.name)
    por_r = {}
    for m in manifest:
        por_r[m["rating"]] = por_r.get(m["rating"], 0) + 1
    print("    por rating: %s" % por_r)
    print("\n  siguiente paso: taggear con WD14 y comparar")
    print("    python Engine/wd14_tagger.py --folder Archivos/wd14/danbooru "
          "--out Archivos/wd14/tags_danbooru")
    print("    python Engine/comparar_danbooru.py")


if __name__ == "__main__":
    main()