"""Extract CLIP features for AF TEST videos (separate from train extraction).

Test-set verifier features are needed ONLY at inference time for the
verifier-as-context experiment (Phase 1). The verifier itself was trained
without test labels (no leakage); this step just runs the trained verifier on
test data so the policy model can see per-second forgery scores as context.
"""
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 TEST_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}", flush=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", flush=True)

    split = "test"
    examples = build_examples(
        annot_dir=ANNOT, video_root=VROOT, generators=TEST_GENERATORS,
        split_prefix=split, preprocessed_data_path=None, require_video_exists=True,
    )
    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
        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()
        labels = frame_labels_from_segments(
            ex["solution"], n_secs, fps_to_groups=SAMPLE_FPS
        ).float()
        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 % 30 == 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()