"""Wan 2.2 LoRA dataset builder.

Feed it a folder of random images/videos; it detects the main subject, crops
to a Wan resolution bucket, captions locally, and writes an
images/ + videos/ + .txt caption dataset plus a starter dataset.toml.

Usage:
    python_embeded\\python.exe dataset_builder\\build.py --input ./raw --output ./dataset_out --trigger "ohwx man"

Run from anywhere -- this script adds its own directory to sys.path so the
sibling modules (crop_resize, detect, caption, scene_split) import cleanly
whether invoked directly or via -m.
"""
import argparse
import os
import sys
import traceback

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import cv2
import numpy as np
from PIL import Image

import crop_resize
import scene_split
from caption import build_captioner, with_trigger
from detect import detect_subject_bbox, subdivide_bbox

IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
VIDEO_EXTS = {".mp4", ".mov", ".mkv", ".avi", ".webm"}


def parse_args():
    p = argparse.ArgumentParser(description="Wan 2.2 LoRA dataset auto-builder")
    p.add_argument("--input", required=True, help="Folder of raw images/videos (scanned recursively)")
    p.add_argument("--output", required=True, help="Output dataset folder")
    p.add_argument("--trigger", default="", help="Trigger word/phrase prepended to every caption")

    p.add_argument("--captioner", choices=["joycaption", "florence2", "wd14"], default="joycaption")
    p.add_argument("--device", default="cuda")
    p.add_argument("--load-in-4bit", action="store_true", help="4-bit load for JoyCaption (bitsandbytes)")
    p.add_argument("--caption-prompt", default=None, help="Override JoyCaption instruction text")
    p.add_argument("--wd14-general-thresh", type=float, default=0.35)
    p.add_argument("--wd14-character-thresh", type=float, default=0.75)

    p.add_argument("--max-tier", choices=["480p", "720p", "1080p"], default="1080p",
                    help="Highest resolution bucket to consider; auto-picks smaller tiers to avoid upscaling")
    p.add_argument("--face-pad", type=float, default=2.2)
    p.add_argument("--person-pad", type=float, default=1.3)
    p.add_argument("--person-region-pad", type=float, default=1.7,
                    help="padding for person_upper/person_lower crops (see --body-region)")
    p.add_argument("--body-region", choices=["auto", "upper", "lower"], default="auto",
                    help="for person-kind detections, narrow the crop to the top or bottom "
                         "portion of the detected body instead of the whole figure -- use "
                         "when most source frames are already close-ups of one body region. "
                         "'auto' (default) keeps the full detected box. No effect on face-kind "
                         "detections.")
    p.add_argument("--body-region-frac", type=float, default=0.45,
                    help="fraction of the person bbox's height kept when --body-region is upper/lower")
    p.add_argument("--min-face-conf", type=float, default=0.5)
    p.add_argument("--min-person-conf", type=float, default=0.35)
    p.add_argument("--image-format", choices=["jpg", "png"], default="jpg")

    p.add_argument("--video-fps", type=float, default=16.0)
    p.add_argument("--clip-min-sec", type=float, default=2.0)
    p.add_argument("--clip-max-sec", type=float, default=6.0)
    p.add_argument("--scene-diff-thresh", type=float, default=0.45)
    p.add_argument("--clips-per-video", type=int, default=5,
                    help="cap clips extracted per source video, evenly sampled across its "
                         "timeline (e.g. a 5-minute video with many scene cuts won't dump "
                         "dozens of near-duplicate clips into the dataset); pass 0 for no cap")

    return p.parse_args()


def scan_inputs(input_dir):
    images, videos = [], []
    for root, _dirs, files in os.walk(input_dir):
        for name in files:
            ext = os.path.splitext(name)[1].lower()
            path = os.path.join(root, name)
            if ext in IMAGE_EXTS:
                images.append(path)
            elif ext in VIDEO_EXTS:
                videos.append(path)
    return sorted(images), sorted(videos)


def process_image(path, idx, out_dir, args, captioner, pad_overrides):
    pil_img = Image.open(path).convert("RGB")
    bgr = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)

    bbox, kind = detect_subject_bbox(bgr, args.min_face_conf, args.min_person_conf)
    if kind == "person" and args.body_region != "auto":
        bbox = subdivide_bbox(bbox, args.body_region, args.body_region_frac)
        kind = f"person_{args.body_region}"
    crop_box, (tw, th), bucket_name, tier = crop_resize.compute_crop(
        pil_img.width, pil_img.height, bbox, kind, pad_overrides, args.max_tier
    )
    cropped = crop_resize.crop_resize_pil(pil_img, crop_box, tw, th)

    stem = f"{idx:04d}"
    images_dir = os.path.join(out_dir, "images")
    os.makedirs(images_dir, exist_ok=True)
    img_path = os.path.join(images_dir, f"{stem}.{args.image_format}")
    if args.image_format == "jpg":
        cropped.save(img_path, quality=95)
    else:
        cropped.save(img_path)

    caption = captioner.caption([cropped])
    caption = with_trigger(caption, args.trigger)
    with open(os.path.join(images_dir, f"{stem}.txt"), "w", encoding="utf-8") as f:
        f.write(caption)

    print(f"  [image] {os.path.basename(path)} -> {stem}.{args.image_format} "
          f"({kind}, {bucket_name}/{tier}, {tw}x{th})")


def subsample_ranges(ranges, target):
    """Evenly picks `target` ranges spread across the full list (by index, so
    spread across the video's timeline) instead of just taking the first N --
    keeps whatever diversity of scenes exists rather than favoring the start."""
    if not target or len(ranges) <= target:
        return ranges
    if target <= 1:
        return ranges[len(ranges) // 2 : len(ranges) // 2 + 1]
    step = (len(ranges) - 1) / (target - 1)
    picked_indices = sorted({round(i * step) for i in range(target)})
    return [ranges[i] for i in picked_indices]


def process_video(path, idx_start, out_dir, args, captioner, pad_overrides):
    cuts = scene_split.find_scene_cuts(path, diff_thresh=args.scene_diff_thresh)
    cap = cv2.VideoCapture(path)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    src_fps = cap.get(cv2.CAP_PROP_FPS) or args.video_fps
    cap.release()

    ranges = scene_split.build_clip_ranges(
        total_frames, src_fps, cuts, args.clip_min_sec, args.clip_max_sec
    )
    detected_count = len(ranges)
    ranges = subsample_ranges(ranges, args.clips_per_video)
    if args.clips_per_video and detected_count > len(ranges):
        print(f"  [video] {os.path.basename(path)}: {detected_count} scene(s) detected, "
              f"sampling {len(ranges)} (--clips-per-video {args.clips_per_video})")

    idx = idx_start
    for start, end in ranges:
        result = scene_split.extract_clip(
            path, start, end, args.video_fps, args.max_tier, pad_overrides,
            args.body_region, args.body_region_frac,
        )
        if result is None:
            continue
        frames, (tw, th), representative, motion_tag = result

        stem = f"{idx:04d}"
        videos_dir = os.path.join(out_dir, "videos")
        os.makedirs(videos_dir, exist_ok=True)
        clip_path = os.path.join(videos_dir, f"{stem}.mp4")
        scene_split.write_clip_mp4(frames, clip_path, args.video_fps)

        caption = captioner.caption([representative])
        caption = f"{caption.strip()}, {motion_tag}" if caption.strip() else motion_tag
        caption = with_trigger(caption, args.trigger)
        with open(os.path.join(videos_dir, f"{stem}.txt"), "w", encoding="utf-8") as f:
            f.write(caption)

        print(f"  [video] {os.path.basename(path)} [{start}-{end}] -> {stem}.mp4 "
              f"({len(frames)}f @ {args.video_fps}fps, {tw}x{th}, {motion_tag})")
        idx += 1
    return idx


def write_dataset_toml(out_dir, args, has_images, has_videos):
    """musubi-tuner dataset config (schema verified against docs/dataset_config.md).
    Items are already pre-cropped to one of a handful of Wan resolution buckets
    by this tool, so enable_bucket lets musubi-tuner group them without forcing
    a single resolution; `resolution` below is just the nominal/reference size."""
    ref_w, ref_h = crop_resize.bucket_set(args.max_tier)["landscape"]

    def p(*parts):
        return os.path.join(out_dir, *parts).replace("\\", "/")

    lines = [
        "[general]",
        f"resolution = [{ref_w}, {ref_h}]",
        'caption_extension = ".txt"',
        "batch_size = 1",
        "enable_bucket = true",
        "bucket_no_upscale = true",
        "",
    ]
    if has_images:
        lines += [
            "[[datasets]]",
            f'image_directory = "{p("images")}"',
            f'cache_directory = "{p("images", "cache")}"',
            "num_repeats = 1",
            "",
        ]
    if has_videos:
        max_frames = int(args.video_fps * args.clip_max_sec) + 1
        lines += [
            "[[datasets]]",
            f'video_directory = "{p("videos")}"',
            f'cache_directory = "{p("videos", "cache")}"',
            'frame_extraction = "full"',
            f"max_frames = {max_frames}",
            f"source_fps = {args.video_fps}",
            "num_repeats = 1",
            "",
        ]
    with open(os.path.join(out_dir, "dataset.toml"), "w", encoding="utf-8") as f:
        f.write("\n".join(lines))


def main():
    args = parse_args()
    os.makedirs(os.path.join(args.output, "images"), exist_ok=True)
    os.makedirs(os.path.join(args.output, "videos"), exist_ok=True)

    images, videos = scan_inputs(args.input)
    print(f"Found {len(images)} image(s), {len(videos)} video(s) under {args.input}")
    if not images and not videos:
        print("Nothing to do.")
        return

    pad_overrides = {
        "face": args.face_pad, "person": args.person_pad,
        "person_upper": args.person_region_pad, "person_lower": args.person_region_pad,
    }

    print(f"Loading captioner: {args.captioner} ...")
    captioner = build_captioner(args)

    img_idx = 1
    for path in images:
        try:
            process_image(path, img_idx, args.output, args, captioner, pad_overrides)
            img_idx += 1
        except Exception:
            print(f"  [image] FAILED {path}:")
            traceback.print_exc()

    vid_idx = 1
    for path in videos:
        try:
            vid_idx = process_video(path, vid_idx, args.output, args, captioner, pad_overrides)
        except Exception:
            print(f"  [video] FAILED {path}:")
            traceback.print_exc()

    write_dataset_toml(args.output, args, bool(images), bool(videos))
    print(f"Done. {img_idx - 1} image(s), {vid_idx - 1} clip(s) -> {args.output}")


if __name__ == "__main__":
    main()