#!/usr/bin/env python3
"""frame_analysis.py — 抽帧 + 豆包视觉逐帧点评.

用法:
  python3 frame_analysis.py --workdir /tmp/lapian --fps 0.1 --max-frames 12
  python3 frame_analysis.py --workdir /tmp/lapian --skip-vision   # 只抽帧, 不调视觉
前置: workdir/video.mp4 (可选 workdir/transcript.json 用于对齐台词)
产物: workdir/frames/f0001.jpg... + workdir/frames.json
  frames.json = [{frame, file, t_ms, t_label, vision, caption}...]
环境: FEME_NEWAPI_TOKEN — 调 api.fmode.cn 视觉模型(计费)
"""
import argparse
import base64
import json
import os
import subprocess
import sys
import urllib.request

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from common import ffprobe_duration, fmt_ts, get_api_token  # noqa: E402

VISION_MODEL = os.environ.get("FEME_VISION_MODEL", "doubao-seed-2-0-pro-260215")
VISION_BASE = os.environ.get("FEME_API_BASE", "https://api.fmode.cn/v1")

PROMPT = """你是短视频编导顾问, 对视频抽帧做画面拆解。用中文 JSON 输出(不要markdown代码块):
{"shot":"景别(远景/全景/中景/近景/特写)","scene":"场景与置景一句话",
 "subjects":"人物/主体在做什么(表情动作)","text_on_screen":"画面文字, 没有则空串",
 "aesthetic":"画面技法一句话(光影/构图/节奏)"},
只描述这一帧可见的信息。"""


def extract_frames(video: str, outdir: str, fps: float, max_frames: int) -> list:
    """均匀抽帧(先抽满 fps, 超出上限再等距丢弃), 返回 [{frame,file,t_ms}]。"""
    os.makedirs(outdir, exist_ok=True)
    subprocess.run(["ffmpeg", "-y", "-i", video, "-vf", f"fps={fps}",
                    "-q:v", "3", os.path.join(outdir, "f%04d.jpg")],
                   capture_output=True, text=True)
    files = sorted(f for f in os.listdir(outdir) if f.endswith(".jpg"))
    if len(files) > max_frames:
        keep = [files[round(i * (len(files) - 1) / (max_frames - 1))]
                for i in range(max_frames)]
        for f in files:
            if f not in keep:
                os.remove(os.path.join(outdir, f))
        files = keep
    dur = ffprobe_duration(video) or 0.0
    frames = []
    for i, f in enumerate(files):
        t_ms = int(dur * 1000 * i / max(len(files) - 1, 1))
        frames.append({"frame": i + 1, "file": f, "t_ms": t_ms,
                       "t_label": fmt_ts(t_ms)})
    return frames


def nearest_line(t_ms: int, segments: list) -> str:
    for seg in segments:
        if seg.get("bg", 0) <= t_ms < seg.get("ed", 0):
            return seg.get("text", "")
    return ""


def call_vision(image_path: str) -> str:
    """单帧 → 豆包视觉, 返回模型文本(JSON字符串或自然语言)。"""
    b64 = base64.b64encode(open(image_path, "rb").read()).decode()
    payload = {
        "model": VISION_MODEL,
        "messages": [{"role": "user", "content": [
            {"type": "text", "text": PROMPT},
            {"type": "image_url",
             "image_url": {"url": f"data:image/jpeg;base64,{b64}"}}]}],
        "temperature": 0.12, "max_tokens": 500,
    }
    req = urllib.request.Request(
        f"{VISION_BASE}/chat/completions",
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json",
                 "Authorization": f"Bearer {get_api_token()}"})
    with urllib.request.urlopen(req, timeout=120) as r:
        resp = json.loads(r.read())
    return (resp["choices"][0]["message"]["content"] or "").strip()


def parse_vision_json(text: str) -> dict:
    """宽容解析模型输出(剥 markdown 围栏, 抠第一个 {...})。"""
    text = text.strip()
    if text.startswith("```"):
        text = text.strip("`").lstrip("json").strip()
    try:
        return json.loads(text)
    except Exception:
        pass
    if "{" in text and "}" in text:
        try:
            return json.loads(text[text.index("{"): text.rindex("}") + 1])
        except Exception:
            pass
    return {}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--workdir", required=True)
    ap.add_argument("--fps", type=float, default=0.15, help="抽帧密度(默认0.15≈6.7s一帧)")
    ap.add_argument("--max-frames", type=int, default=12, help="最多分析帧数(控制计费)")
    ap.add_argument("--skip-vision", action="store_true", help="只抽帧, 不调视觉模型")
    args = ap.parse_args()

    video = os.path.join(args.workdir, "video.mp4")
    if not os.path.exists(video):
        raise SystemExit(f"[frames] 缺少 {video}")
    outdir = os.path.join(args.workdir, "frames")
    frames = extract_frames(video, outdir, args.fps, args.max_frames)
    print(f"[frames] 抽帧 {len(frames)} 张 → {outdir}")

    segs = []
    tp = os.path.join(args.workdir, "transcript.json")
    if os.path.exists(tp):
        segs = (json.load(open(tp, encoding="utf-8")).get("data") or {}).get("segments") or []

    for fr in frames:
        fr["line"] = nearest_line(fr["t_ms"], segs)
        if args.skip_vision:
            fr["vision"] = {}
            fr["caption"] = fr["line"][:60] or "(未调视觉)"
            continue
        img = os.path.join(outdir, fr["file"])
        try:
            raw = call_vision(img)
            v = parse_vision_json(raw)
            fr["vision"] = v
            parts = [v.get(k, "") for k in ("scene", "subjects", "aesthetic")]
            fr["caption"] = "；".join(p for p in parts if p) or raw[:120]
            print(f"[frames] f{fr['frame']:04d} {fr['t_label']} ✓ {fr['caption'][:50]}")
        except Exception as e:
            fr["vision"], fr["caption"] = {}, f"(视觉识别失败: {e})"
            print(f"[frames] f{fr['frame']:04d} {fr['t_label']} ✗ {e}")

    out = os.path.join(args.workdir, "frames.json")
    json.dump(frames, open(out, "w", encoding="utf-8"), ensure_ascii=False, indent=2)
    ok = sum(1 for f in frames if f.get("vision"))
    print(f"[frames] frames.json: {len(frames)} 帧(视觉成功 {ok})")


if __name__ == "__main__":
    main()