#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""从 annotations.jsonl 生成 labels.csv（帧路径, label, task）。"""
import argparse
import csv
import json
import os
import random
import subprocess
import sys
from pathlib import Path

import yaml

ROOT = Path(__file__).resolve().parent.parent


def load_config(path):
    with open(path, encoding="utf-8") as f:
        cfg = yaml.safe_load(f)
    base = Path(path).resolve().parent
    for key in ("annotations_jsonl", "frames_dir", "labels_csv"):
        p = cfg["data"][key]
        if not os.path.isabs(p):
            cfg["data"][key] = str((base / p).resolve())
    cfg["data"]["output_dir"] = cfg["export"]["output_dir"]
    if not os.path.isabs(cfg["export"]["output_dir"]):
        cfg["export"]["output_dir"] = str((base / cfg["export"]["output_dir"]).resolve())
    return cfg


def ffprobe_duration(video_path, ffmpeg_dir=""):
    ffprobe = "ffprobe"
    if ffmpeg_dir:
        ffprobe = os.path.join(ffmpeg_dir, "ffprobe.exe" if os.name == "nt" else "ffprobe")
    cmd = [
        ffprobe, "-v", "error", "-show_entries", "format=duration",
        "-of", "default=noprint_wrappers=1:nokey=1", video_path,
    ]
    try:
        out = subprocess.check_output(cmd, stderr=subprocess.DEVNULL, text=True).strip()
        return float(out) if out else 0.0
    except Exception:
        return 0.0


def extract_frame_at(video_path, out_path, sec, ffmpeg_dir=""):
    ffmpeg = "ffmpeg"
    if ffmpeg_dir:
        ffmpeg = os.path.join(ffmpeg_dir, "ffmpeg.exe" if os.name == "nt" else "ffmpeg")
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    cmd = [
        ffmpeg, "-y", "-ss", str(sec), "-i", video_path,
        "-frames:v", "1", "-q:v", "2", out_path,
    ]
    subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=False)


def process_video(item, cfg, ffmpeg_dir):
    media_id = item["media_id"]
    video_path = item["video_path"]
    if not os.path.isfile(video_path) and not video_path.startswith("http"):
        print(f"跳过 media_id={media_id}: 视频不存在 {video_path}")
        return []

    local_video = video_path
    tmp_video = None
    if video_path.startswith("http"):
        import httpx
        tmp_video = os.path.join(cfg["data"]["frames_dir"], f"_tmp_{media_id}.mp4")
        os.makedirs(cfg["data"]["frames_dir"], exist_ok=True)
        with httpx.stream("GET", video_path, timeout=600.0) as r:
            r.raise_for_status()
            with open(tmp_video, "wb") as f:
                for chunk in r.iter_bytes():
                    f.write(chunk)
        local_video = tmp_video

    duration = ffprobe_duration(local_video, ffmpeg_dir)
    if duration <= 0:
        duration = max(item.get("handover_time_sec", 600), 600)

    sample_fps = cfg["sampling"]["sample_fps"]
    win = cfg["sampling"]["positive_window_sec"]
    sf_t = float(item["storefront_time_sec"])
    ho_t = float(item["handover_time_sec"])

    rows = []
    t = 0.0
    step = 1.0 / sample_fps
    other_count = 0
    other_budget = max(1, int(duration * sample_fps / cfg["sampling"]["other_downsample_ratio"]))
    while t <= duration:
        rel = f"{media_id}/{int(t * 1000)}.jpg"
        out_path = os.path.join(cfg["data"]["frames_dir"], rel)
        is_sf = abs(t - sf_t) <= win
        is_ho = abs(t - ho_t) <= win
        need_other = not is_sf and not is_ho and other_count < other_budget
        if is_sf or is_ho or need_other:
            extract_frame_at(local_video, out_path, t, ffmpeg_dir)
            if os.path.isfile(out_path):
                split = item.get("split", "train")
                sf_label = 1 if is_sf else 0
                ho_label = 1 if is_ho else 0
                rows.append((rel, sf_label, "storefront", split))
                rows.append((rel, ho_label, "handover", split))
                if need_other:
                    other_count += 1
        t += step
    if tmp_video and os.path.isfile(tmp_video):
        os.remove(tmp_video)
    return rows


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--config", default=str(Path(__file__).parent / "config.yaml"))
    parser.add_argument("--ffmpeg-dir", default=os.environ.get("FFMPEG_DIR", ""))
    parser.add_argument("--limit", type=int, default=0)
    args = parser.parse_args()
    cfg = load_config(args.config)
    os.makedirs(cfg["data"]["frames_dir"], exist_ok=True)

    items = []
    with open(cfg["data"]["annotations_jsonl"], encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                items.append(json.loads(line))
    if args.limit:
        items = items[: args.limit]

    all_rows = []
    for i, item in enumerate(items):
        if item.get("storefront_time_sec", 0) <= 0 or item.get("handover_time_sec", 0) <= 0:
            print(f"跳过未标注 media_id={item.get('media_id')}")
            continue
        print(f"[{i+1}/{len(items)}] media_id={item['media_id']}")
        all_rows.extend(process_video(item, cfg, args.ffmpeg_dir))

    with open(cfg["data"]["labels_csv"], "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["frame_path", "label", "task", "split"])
        for row in all_rows:
            w.writerow(row)
    print(f"写入 {len(all_rows)} 行 -> {cfg['data']['labels_csv']}")


if __name__ == "__main__":
    main()