import argparse
import json
from pathlib import Path

import torch
from PIL import Image
from transformers import AutoImageProcessor, VideoMAEForVideoClassification


def clamp01(value):
    try:
        n = float(value)
    except Exception:
        return 0.0
    return max(0.0, min(1.0, n))


def existing_model_dir(path: Path):
    try:
        if path.exists() and path.is_dir() and (path / "config.json").exists():
            return path
    except Exception:
        pass
    return None


def resolve_model_path(root: Path, requested: str):
    requested = str(requested or "").strip()
    if requested:
        p = existing_model_dir(Path(requested).expanduser().resolve())
        if p:
            return p, "videomae_model"
        return Path(requested).expanduser(), "videomae_missing"

    trained = root / "videomae" / "model"
    p = existing_model_dir(trained)
    if p:
        return p, "videomae_clip"

    return trained, "videomae_missing"


def resample_values(values: list, count: int) -> list:
    if not values:
        return []
    if len(values) == 1:
        return [values[0] for _ in range(count)]
    if count <= 1:
        return [values[0]]

    last = len(values) - 1
    return [values[int(round((i * last) / max(1, count - 1)))] for i in range(count)]


def load_frames(paths: list[str], num_frames: int):
    selected = resample_values(paths, num_frames)
    frames = []
    for path in selected:
        with Image.open(path) as img:
            frames.append(img.convert("RGB").copy())
    return frames


def frame_paths_from_args(args):
    paths = []
    if args.frames_json:
        with Path(args.frames_json).open("r", encoding="utf-8") as f:
            data = json.load(f)
        paths.extend(str(p) for p in data)

    if args.clip_dir:
        clip_dir = Path(args.clip_dir)
        paths.extend(
            str(p) for p in sorted(clip_dir.iterdir())
            if p.is_file() and p.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp"}
        )

    paths.extend(str(p) for p in args.frames)
    return [p for p in paths if p.strip()]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", required=True)
    parser.add_argument("--model", default="")
    parser.add_argument("--clip-dir", default="")
    parser.add_argument("--frames-json", default="")
    parser.add_argument("--frames", nargs="*", default=[])
    parser.add_argument("--num-frames", type=int, default=16)
    parser.add_argument("--device", default="auto")
    args = parser.parse_args()

    root = Path(args.root).resolve()
    model_path, model_source = resolve_model_path(root, args.model)
    if not existing_model_dir(model_path):
        print(json.dumps({
            "available": False,
            "source": model_source,
            "modelPath": str(model_path),
            "sexPosition": "keine",
            "sexPositionScore": 0.0,
            "scores": [],
        }, ensure_ascii=False))
        return

    paths = frame_paths_from_args(args)
    if not paths:
        print(json.dumps({
            "available": False,
            "source": "videomae_no_frames",
            "modelPath": str(model_path),
            "sexPosition": "keine",
            "sexPositionScore": 0.0,
            "scores": [],
        }, ensure_ascii=False))
        return

    if str(args.device).lower() == "auto":
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        device = torch.device(args.device)

    try:
        processor = AutoImageProcessor.from_pretrained(model_path)
        model = VideoMAEForVideoClassification.from_pretrained(model_path).to(device)
        model.eval()

        frames = load_frames(paths, max(2, int(args.num_frames or 16)))
        inputs = processor(frames, return_tensors="pt")
        pixel_values = inputs["pixel_values"].to(device)

        with torch.no_grad():
            logits = model(pixel_values=pixel_values).logits
            probs = torch.softmax(logits, dim=-1)[0].detach().cpu().tolist()

        id_to_label = getattr(model.config, "id2label", {}) or {}
        scores = []
        best_label = "keine"
        best_score = 0.0

        for idx, score in enumerate(probs):
            label = str(id_to_label.get(idx, id_to_label.get(str(idx), idx))).strip()
            score = clamp01(score)
            scores.append({"label": label, "score": score})
            if score > best_score:
                best_label = label
                best_score = score

        scores.sort(key=lambda item: item["score"], reverse=True)

        print(json.dumps({
            "available": True,
            "source": model_source,
            "modelPath": str(model_path),
            "device": str(device),
            "sexPosition": best_label,
            "sexPositionScore": best_score,
            "scores": scores[:10],
        }, ensure_ascii=False))

    except Exception as exc:
        print(json.dumps({
            "available": False,
            "source": "videomae_predict_failed",
            "modelPath": str(model_path),
            "error": repr(exc),
            "sexPosition": "keine",
            "sexPositionScore": 0.0,
            "scores": [],
        }, ensure_ascii=False))


if __name__ == "__main__":
    main()