"""M2 step 1: extract per-second CLIP ViT-L/14 features for AF TRAIN videos.

Output layout (matches existing preprocess cache style):
  <CACHE>/train/<gen>/<sample_id>/clip_feats.pt   (T, 768) float32
  <CACHE>/train/<gen>/<sample_id>/clip_labels.pt  (T,)    float32 binary

Sampling: 1 frame / source-second (matches our fps_to_groups=1.0 convention).

TEST SPLIT IS DELIBERATELY NOT EXTRACTED.  Test videos are held out for final
policy evaluation; using them anywhere in verifier training (even for feature
caching with leaked side-effects) would compromise the experimental setup.
Test extraction is a separate script run only at final benchmarking.
"""
import argparse
import os
import sys
import time

import torch
import torch.nn.functional as F
from decord import VideoReader, cpu
from PIL import Image
from transformers import CLIPModel, CLIPProcessor

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from src.open_r1.data_loader import TRAIN_GENERATORS, build_examples
from src.open_r1.forgery_head import frame_labels_from_segments

ANNOT = "/mnt/local-fast/zhangt/annot/annot"
VROOT = "/mnt/local-fast/zhangt/video"
CACHE = "/mnt/local-fast/zhangt/forensics_verifier_clip_l14"
MODEL_ID = "openai/clip-vit-large-patch14"
SAMPLE_FPS = 1.0


def decode_at_1fps(video_path: str, duration: float):
    vr = VideoReader(video_path, ctx=cpu(0))
    fps_video = vr.get_avg_fps()
    n_secs = max(1, int(duration))
    idxs = [min(int(s * fps_video), len(vr) - 1) for s in range(n_secs)]
    frames = vr.get_batch(idxs).asnumpy()
    return [Image.fromarray(f) for f in frames], n_secs


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--rank", type=int, default=0)
    ap.add_argument("--world_size", type=int, default=1)
    ap.add_argument("--device", type=int, default=0)
    ap.add_argument("--batch", type=int, default=32)
    args = ap.parse_args()

    device = f"cuda:{args.device}"
    print(f"[rank {args.rank}/{args.world_size}] device={device}  cache={CACHE}", flush=True)
    os.makedirs(CACHE, exist_ok=True)

    print("loading CLIP ...", flush=True)
    t0 = time.time()
    clip = CLIPModel.from_pretrained(MODEL_ID, torch_dtype=torch.float32).to(device).eval()
    proc = CLIPProcessor.from_pretrained(MODEL_ID)
    print(f"  loaded in {time.time()-t0:.1f}s. proj_dim={clip.config.projection_dim}", flush=True)

    split = "train"
    examples = build_examples(
        annot_dir=ANNOT, video_root=VROOT, generators=TRAIN_GENERATORS,
        split_prefix=split, preprocessed_data_path=None, require_video_exists=True,
    )
    # Shard for parallel run
    examples = [ex for i, ex in enumerate(examples) if i % args.world_size == args.rank]
    print(f"[{split}] rank={args.rank} processing {len(examples)} videos", flush=True)

    t_start = time.time()
    done = skipped = failed = 0
    for ex in examples:
        sample_id = os.path.splitext(os.path.basename(ex["video_path"]))[0]
        out_dir = os.path.join(CACHE, split, ex["generator"], sample_id)
        feat_path = os.path.join(out_dir, "clip_feats.pt")
        if os.path.exists(feat_path):
            skipped += 1
            continue
        try:
            pil_imgs, n_secs = decode_at_1fps(ex["video_path"], ex["durations"])
        except Exception:
            failed += 1
            continue

        # Batched CLIP forward
        feats_all = []
        for i in range(0, len(pil_imgs), args.batch):
            chunk = pil_imgs[i:i + args.batch]
            with torch.no_grad():
                inputs = proc(images=chunk, return_tensors="pt").to(device)
                f = clip.get_image_features(**inputs)
                f = F.normalize(f, dim=-1)
            feats_all.append(f.cpu())
        feats = torch.cat(feats_all, dim=0).float()       # (T, 768)
        labels = frame_labels_from_segments(
            ex["solution"], n_secs, fps_to_groups=SAMPLE_FPS
        ).float()                                          # (T,)
        assert feats.shape[0] == labels.shape[0], (feats.shape, labels.shape)

        os.makedirs(out_dir, exist_ok=True)
        torch.save(feats, feat_path)
        torch.save(labels, os.path.join(out_dir, "clip_labels.pt"))
        done += 1

        if done % 50 == 0:
            elapsed = time.time() - t_start
            rate = done / max(1e-6, elapsed)
            remaining = (len(examples) - done - skipped) / max(1e-6, rate)
            print(
                f"  rank={args.rank} done={done} skipped={skipped} failed={failed} "
                f"rate={rate:.2f}/s eta={remaining/60:.1f}min",
                flush=True,
            )

    print(f"rank={args.rank} DONE  done={done} skipped={skipped} failed={failed}", flush=True)


if __name__ == "__main__":
    main()