#!/usr/bin/env python3
"""WD14 tagger — danbooru-теги для аниме-картинок (SmilingWolf/wd-*-tagger-v3).

Пишет теги в <имя_файла>.txt (danbooru-стиль: "1girl, animal ears, ...").
Уже размеченные файлы пропускает (--force — перезаписать).

Usage:
    python wd14_tagger.py <папка_или_файл>... [--model vit] [--gen 0.5] [--char 0.85]

Зависимости: torch, timm, pandas, pillow, huggingface_hub. Нужен GPU (или CPU).
"""

import argparse
import time
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pandas as pd
import timm
import torch
from huggingface_hub import hf_hub_download
from huggingface_hub.utils import HfHubHTTPError
from PIL import Image
from timm.data import create_transform, resolve_data_config
from torch import nn
from torch.nn import functional as F
from tqdm import tqdm

MODEL_REPO_MAP = {
    "vit": "SmilingWolf/wd-vit-tagger-v3",
    "swinv2": "SmilingWolf/wd-swinv2-tagger-v3",
    "convnext": "SmilingWolf/wd-convnext-tagger-v3",
}

IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}


def pil_ensure_rgb(image: Image.Image) -> Image.Image:
    if image.mode not in ("RGB", "RGBA"):
        image = image.convert("RGBA" if "transparency" in image.info else "RGB")
    if image.mode == "RGBA":
        canvas = Image.new("RGBA", image.size, (255, 255, 255))
        canvas.alpha_composite(image)
        image = canvas.convert("RGB")
    return image


def pil_pad_square(image: Image.Image) -> Image.Image:
    w, h = image.size
    px = max(image.size)
    canvas = Image.new("RGB", (px, px), (255, 255, 255))
    canvas.paste(image, ((px - w) // 2, (px - h) // 2))
    return canvas


@dataclass
class LabelData:
    names: list
    rating: list
    general: list
    character: list


def load_labels_hf(repo_id: str) -> LabelData:
    try:
        csv_path = hf_hub_download(repo_id=repo_id, filename="selected_tags.csv")
    except HfHubHTTPError as e:
        raise FileNotFoundError(f"selected_tags.csv failed to download from {repo_id}") from e
    df = pd.read_csv(csv_path, usecols=["name", "category"])
    return LabelData(
        names=df["name"].tolist(),
        rating=list(np.where(df["category"] == 9)[0]),
        general=list(np.where(df["category"] == 0)[0]),
        character=list(np.where(df["category"] == 4)[0]),
    )


def get_tags(probs, labels, gen_threshold, char_threshold):
    """Probs -> строка тегов (danbooru-стиль, рейтинг в конце)."""
    probs = list(zip(labels.names, probs.numpy()))
    rating = dict(sorted((probs[i] for i in labels.rating), key=lambda x: x[1], reverse=True))
    gen = {n: p for n, p in (probs[i] for i in labels.general) if p > gen_threshold}
    char = {n: p for n, p in (probs[i] for i in labels.character) if p > char_threshold}
    gen = dict(sorted(gen.items(), key=lambda x: x[1], reverse=True))
    char = dict(sorted(char.items(), key=lambda x: x[1], reverse=True))
    tags = list(gen) + list(char)
    caption = ", ".join(t.replace("_", " ").replace("(", r"\(").replace(")", r"\)") for t in tags)
    caption += ", rating_" + next(iter(rating), "")
    return caption


def main():
    parser = argparse.ArgumentParser(description="WD14 tagger (danbooru-теги)")
    parser.add_argument("paths", nargs="+", help="папки или файлы (папки обходятся рекурсивно)")
    parser.add_argument("--model", default="vit", choices=sorted(MODEL_REPO_MAP))
    parser.add_argument("--gen", dest="gen_threshold", type=float, default=0.5,
                        help="порог уверенности для general-тегов")
    parser.add_argument("--char", dest="char_threshold", type=float, default=0.85,
                        help="порог уверенности для character-тегов")
    parser.add_argument("--force", action="store_true", help="перезаписать существующие .txt")
    args = parser.parse_args()

    repo_id = MODEL_REPO_MAP[args.model]
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # собираем файлы
    files = []
    for arg in args.paths:
        p = Path(arg)
        if p.is_dir():
            files.extend(f for f in p.rglob("*") if f.suffix.lower() in IMAGE_EXTS)
        elif p.is_file():
            files.append(p)
    files = sorted(set(files))
    if not files:
        print("Нет файлов для обработки.")
        return

    print(f"Загрузка модели {args.model} ({repo_id}) на {device}...")
    model = timm.create_model("hf-hub:" + repo_id).eval()
    model.load_state_dict(timm.models.load_state_dict_from_hf(repo_id))
    labels = load_labels_hf(repo_id)
    transform = create_transform(**resolve_data_config(model.pretrained_cfg, model=model))
    if device.type == "cuda":
        model = model.to(device)

    t0 = time.time()
    done = skipped = 0
    for image_path in tqdm(files, desc="Tagging", unit="img"):
        txt_path = image_path.with_suffix(".txt")
        if txt_path.exists() and not args.force:
            skipped += 1
            continue
        try:
            img = pil_pad_square(pil_ensure_rgb(Image.open(image_path)))
            inputs = transform(img).unsqueeze(0)[:, [2, 1, 0]]  # RGB -> BGR
            with torch.inference_mode():
                outputs = F.sigmoid(model(inputs.to(device))).cpu()
            caption = get_tags(outputs.squeeze(0), labels, args.gen_threshold, args.char_threshold)
            txt_path.write_text(caption + "\n")
            done += 1
        except Exception as e:
            tqdm.write(f"  SKIP {image_path.name}: {e}")

    print(f"\nГотово: {done} размечено, {skipped} пропущено (уже есть .txt), "
          f"за {time.strftime('%H:%M:%S', time.gmtime(time.time() - t0))}")


if __name__ == "__main__":
    main()