#!/usr/bin/env python3
"""anime-facesaver: save one still per scene from videos that contain an
anime face.

Fork of facesaver (https://github.com/.../facesaver) fixing its core bug:
the original used stock COCO yolov11l with class 0 ('person') as a proxy
for 'face', so it happily saved stills of headless bodies, hands and
shoes. This version uses a dedicated anime-face detector
(yolov8x6_animeface, AGPL-3.0, by Fuyucchi) and only saves frames where
an actual face of sufficient size is found.

Scene-change gate preserved from the original: at most one still is saved
per detected scene, which keeps the output dataset from exploding with
near-duplicates.
"""

import argparse
import os
import sys


def parse_arguments():
    parser = argparse.ArgumentParser(
        description=(
            "Detect anime faces in videos and save one still per scene "
            "where a qualifying face appears."
        ),
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--input-dir", "-I",
        required=True,
        help="Directory containing input video files.",
    )
    parser.add_argument(
        "--output-dir", "-O",
        required=True,
        help="Directory to save screenshot outputs.",
    )
    parser.add_argument(
        "--gpu", "-g",
        default="1",
        help=(
            "GPU index to pin to (e.g. '1'), or 'cpu' to force CPU. "
            "Use 'auto' to let torch pick the first visible device. "
            "On the shared .80 box, GPU0 is usually held by the inference "
            "server — pinning to GPU1 avoids OOM."
        ),
    )
    parser.add_argument(
        "--model",
        default=os.path.join(os.path.dirname(os.path.abspath(__file__)),
                             "models", "yolov8x6_animeface.pt"),
        help="Path to the anime-face YOLO weights.",
    )
    parser.add_argument(
        "--download-model",
        action="store_true",
        help="Download the yolov8x6_animeface weights if --model is missing.",
    )
    parser.add_argument(
        "--min-width", "-w",
        type=int,
        default=100,
        help="Minimum face bounding box width (px) to trigger a save.",
    )
    parser.add_argument(
        "--min-height", "-m",
        type=int,
        default=100,
        help="Minimum face bounding box height (px) to trigger a save.",
    )
    parser.add_argument(
        "--min-face-fraction",
        type=float,
        default=0.005,
        help="Minimum face box area as a fraction of frame area (kills "
             "background-extra faces in wide shots; 0 disables).",
    )
    parser.add_argument(
        "--max-face-fraction",
        type=float,
        default=0.30,
        help="Maximum face box area as a fraction of frame area (kills "
             "detector meltdowns on full-frame effects like ED feather "
             "storms; 0 disables).",
    )
    parser.add_argument(
        "--min-aspect",
        type=float,
        default=0.45,
        help="Minimum face box width/height ratio.",
    )
    parser.add_argument(
        "--max-aspect",
        type=float,
        default=1.6,
        help="Maximum face box width/height ratio (a face wider than 1.6x "
             "itself is almost always an effect blob).",
    )
    parser.add_argument(
        "--conf",
        type=float,
        default=0.3,
        help="Detector confidence threshold (model ships recommending 0.3; "
             "it is high-precision, raise for a stricter dataset).",
    )
    parser.add_argument(
        "--iou",
        type=float,
        default=0.5,
        help="NMS IoU threshold.",
    )
    parser.add_argument(
        "--edge-margin",
        type=int,
        default=0,
        help="Require the face box to stay this many px inside all frame "
             "edges. 0 (default) allows cropped faces; set >0 to reject "
             "faces touching the crop lines.",
    )
    parser.add_argument(
        "--scene-threshold",
        type=float,
        default=30.0,
        help="PySceneDetect ContentDetector threshold (same as original).",
    )
    parser.add_argument(
        "--imgsz",
        type=int,
        default=1280,
        help="Inference image size. The model was trained at 1280; lower "
             "it (640) for speed at some recall cost.",
    )
    parser.add_argument(
        "--preview-dir",
        default=None,
        help="Optional: also copy each SAVED frame plus an annotated "
             "version (face boxes drawn) here for eyeball QC.",
    )
    return parser.parse_args()


def resolve_gpu(gpu_arg):
    """Pin CUDA before torch initializes. Returns torch device string.

    NOTE: ultralytics' select_device() rewrites CUDA_VISIBLE_DEVICES on
    first predict() if the CUDA context is not yet initialized, silently
    un-pinning us onto GPU0. We force-initialize a tiny CUDA allocation
    here so our CVD pin is locked in before ultralytics can touch it.
    """
    if gpu_arg.lower() == "cpu":
        os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
        return "cpu"
    if gpu_arg.lower() == "auto":
        return None  # let ultralytics/torch decide
    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_arg
    import torch
    if torch.cuda.is_available():
        torch.zeros(1, device="cuda")  # lock the context to our pin
    else:
        return "cpu"
    return "cuda:0"  # within our CVD, physical GPU N is always cuda:0


def ensure_directory(directory):
    if directory and not os.path.exists(directory):
        os.makedirs(directory)


def download_model(path):
    """Fetch yolov8x6_animeface.pt from Hugging Face (Fuyucchi/yolov8_animeface)."""
    url = ("https://huggingface.co/Fuyucchi/yolov8_animeface/"
           "resolve/main/yolov8x6_animeface.pt")
    ensure_directory(os.path.dirname(path))
    tmp = path + ".part"
    print(f"Downloading anime-face model to {path} ...")
    import urllib.request

    def _hook(n, bs, total):
        if total > 0 and n % 20000 == 0:
            pct = min(100.0, n * bs * 100.0 / total)
            print(f"\r  {pct:5.1f}%", end="", flush=True)

    urllib.request.urlretrieve(url, tmp, reporthook=_hook)
    print()
    os.replace(tmp, path)
    print(f"Saved model ({os.path.getsize(path) / 1e6:.0f} MB).")


def face_qualifies(box, frame_shape, min_width, min_height,
                   min_face_fraction, max_face_fraction,
                   min_aspect, max_aspect, edge_margin):
    """Original is_full_face(), corrected: the box must actually be a face
    (enforced by the detector class), size gates stay, hard
    'must-not-touch-edge' rule becomes an optional margin.
    Upper-bound gates (max area, aspect) kill detector meltdowns on
    full-frame stylized effects (ED feather storms etc.)."""
    x1, y1, x2, y2 = box
    frame_h, frame_w = frame_shape[:2]

    w, h = x2 - x1, y2 - y1
    if w < min_width or h < min_height:
        return False

    area_frac = (w * h) / float(frame_w * frame_h)
    if min_face_fraction > 0 and area_frac < min_face_fraction:
        return False
    if max_face_fraction > 0 and area_frac > max_face_fraction:
        return False

    if h > 0:
        aspect = w / float(h)
        if aspect < min_aspect or aspect > max_aspect:
            return False

    if edge_margin > 0:
        if (x1 < edge_margin or y1 < edge_margin
                or x2 > frame_w - edge_margin
                or y2 > frame_h - edge_margin):
            return False

    return True


def process_video(video_path, output_dir, args, model, device, draw,
                  output_start=0):
    from scenedetect import open_video, SceneManager, ContentDetector

    video_name = os.path.splitext(os.path.basename(video_path))[0]

    try:
        import cv2
        video = open_video(video_path)
        scene_manager = SceneManager()
        scene_manager.add_detector(
            ContentDetector(threshold=args.scene_threshold))
    except Exception as e:
        print(f"Error initializing video for scene detection in "
              f"{video_path}: {e}")
        return output_start

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"Error opening video file {video_path}")
        return output_start

    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    try:
        scene_manager.detect_scenes(video=video)
        scene_list = scene_manager.get_scene_list()
        scene_starts = {scene[0].get_frames() for scene in scene_list}
    except Exception as e:
        print(f"Error detecting scenes in {video_path}: {e}")
        cap.release()
        return output_start

    print(f"{video_name}: {frame_count} frames, "
          f"{len(scene_starts)} scene starts")

    output_count = output_start
    face_in_scene = False
    frame_idx = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        if frame_idx in scene_starts:
            face_in_scene = False
        # membership test on a set; pop matched starts to keep it cheap
        # (scenes strictly increase, so set stays small either way)

        if not face_in_scene:
            try:
                results = model.predict(frame, conf=args.conf,
                                        iou=args.iou, imgsz=args.imgsz,
                                        device=device, verbose=False)
                for result in results:
                    if result.boxes is None:
                        continue
                    boxes = result.boxes.xyxy.cpu().numpy()
                    for box in boxes:
                        if face_qualifies(box, frame.shape, args.min_width,
                                          args.min_height,
                                          args.min_face_fraction,
                                          args.max_face_fraction,
                                          args.min_aspect, args.max_aspect,
                                          args.edge_margin):
                            out = os.path.join(
                                output_dir,
                                f"{video_name}_face_{output_count:04d}.png")
                            cv2.imwrite(out, frame)
                            print(f"  saved {os.path.basename(out)} "
                                  f"(frame {frame_idx}, "
                                  f"box {int(box[2]-box[0])}x"
                                  f"{int(box[3]-box[1])})")
                            output_count += 1
                            face_in_scene = True
                            if args.preview_dir:
                                annotated = draw(frame.copy(),
                                                 result.boxes.xyxy.cpu())
                                cv2.imwrite(
                                    os.path.join(
                                        args.preview_dir,
                                        os.path.basename(out).replace(
                                            ".png", "_ann.png")),
                                    annotated)
                            break
                    if face_in_scene:
                        break
            except Exception as e:
                print(f"Error during face detection in {video_path} "
                      f"at frame {frame_idx}: {e}")

        frame_idx += 1

    cap.release()
    print(f"Processed {video_name}: "
          f"{output_count - output_start} screenshots saved.")
    return output_count


def main():
    args = parse_arguments()

    if not os.path.isdir(args.input_dir):
        print(f"Error: Input directory '{args.input_dir}' does not exist.")
        sys.exit(1)

    if args.download_model and not os.path.exists(args.model):
        download_model(args.model)
    if not os.path.exists(args.model):
        print(f"Error: model weights not found at {args.model}\n"
              f"Re-run with --download-model, or place the weights there "
              f"(see README).")
        sys.exit(1)

    device = resolve_gpu(args.gpu)

    import torch
    from ultralytics import YOLO
    import cv2

    if device == "cuda:0" and not torch.cuda.is_available():
        print("CUDA requested but not visible after pinning "
              f"(--gpu {args.gpu}). Falling back to CPU.")
        device = "cpu"
    print(f"Device: {device}"
          + (f" (physical GPU {args.gpu})" if device == "cuda" else ""))

    model = YOLO(args.model)
    names = model.names
    print(f"Model loaded. Classes: {names}")

    def draw(img, boxes):
        for (a, b, c, d) in boxes.int().tolist():
            cv2.rectangle(img, (a, b), (c, d), (0, 255, 0), 2)
        return img

    ensure_directory(args.output_dir)
    if args.preview_dir:
        ensure_directory(args.preview_dir)

    video_extensions = ('.mp4', '.avi', '.mov', '.mkv')
    grand_total = 0
    for filename in sorted(os.listdir(args.input_dir)):
        if filename.lower().endswith(video_extensions):
            video_path = os.path.join(args.input_dir, filename)
            print(f"Processing video: {video_path}")
            grand_total += process_video(
                video_path, args.output_dir, args, model, device, draw)

    print(f"Done. {grand_total} stills saved to {args.output_dir}")


if __name__ == "__main__":
    main()