Video Breakdown

"""
视频预处理脚本（独立可执行）。
自包含视频预处理，无需外部后端服务。
需要本机安装 FFmpeg。

Usage:
    python scripts/process_video.py "<video_url>"

Env:
    FFMPEG_BIN, FFPROBE_BIN, VOLCENGINE_ACCESS_KEY, VOLCENGINE_SECRET_KEY,
    VOLC_ASR_APP_ID, VOLC_ASR_ACCESS_KEY, VOLC_ASR_RESOURCE_ID,
    DATABASE_TOS_BUCKET (或 TOS_BUCKET), DATABASE_TOS_REGION (或 TOS_REGION)
"""

import asyncio
import json
import os
import shutil
import subprocess
import sys
import tempfile
import uuid
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List

import httpx


def _run_command(cmd: List[str]) -> str:
    process = subprocess.run(
        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True
    )
    return process.stdout.strip()


def _probe_video(ffprobe_bin: str, video_path: Path) -> Dict[str, Any]:
    cmd = [
        ffprobe_bin,
        "-v",
        "error",
        "-print_format",
        "json",
        "-show_format",
        "-show_streams",
        str(video_path),
    ]
    result = _run_command(cmd)
    info = json.loads(result) if result else {}
    fmt = info.get("format") or {}
    duration = float(fmt.get("duration", 0))
    w, h, fr = None, None, None
    for s in info.get("streams", []):
        if s.get("codec_type") == "video":
            w, h, fr = s.get("width"), s.get("height"), s.get("r_frame_rate")
            break
    return {
        "duration": duration,
        "width": w,
        "height": h,
        "frame_rate": fr,
        "size": fmt.get("size"),
    }


def _build_segments(duration: float):
    segments = []
    idx = 1
    bps = [0.0, 3.0, 5.0, 10.0, 20.0]
    for i in range(len(bps) - 1):
        if duration <= bps[i]:
            break
        end = min(bps[i + 1], duration)
        if end - bps[i] < 0.5:
            break
        segments.append({"index": idx, "start": bps[i], "end": end})
        idx += 1
    if duration > 20.0:
        c = 20.0
        while c < duration:
            e = min(duration, c + 10.0)
            if e - c < 0.5:
                break
            segments.append({"index": idx, "start": c, "end": e})
            idx += 1
            c = e
    return segments


def _extract_frames(ffmpeg_bin: str, video_path: Path, segments, fps=3):
    frames_dir = video_path.parent / "frames"
    frames_dir.mkdir(exist_ok=True)
    for seg in segments:
        seg["frame_paths"] = []
        dur = seg["end"] - seg["start"]
        for i in range(fps):
            ratio = i / max(fps - 1, 1)
            offset = min(seg["start"] + ratio * dur, seg["end"] - 0.1)
            offset = max(offset, seg["start"])
            out = frames_dir / f"seg{seg['index']:03d}_frame_{i}.jpg"
            cmd = [
                ffmpeg_bin,
                "-y",
                "-ss",
                f"{offset:.2f}",
                "-i",
                str(video_path),
                "-frames:v",
                "1",
                "-q:v",
                "5",
                str(out),
            ]
            try:
                _run_command(cmd)
                if out.exists():
                    seg["frame_paths"].append(str(out))
            except Exception:
                pass


async def main(video_url: str):
    ffmpeg_bin = os.getenv("FFMPEG_BIN", "ffmpeg")
    ffprobe_bin = os.getenv("FFPROBE_BIN", "ffprobe")
    temp_dir = Path(tempfile.mkdtemp(prefix="vba_"))
    task_id = datetime.now().strftime("%Y%m%d_%H%M%S") + "_" + uuid.uuid4().hex[:8]
    local_video = temp_dir / f"{task_id}.mp4"

    try:
        # Download
        print("[1/4] 下载视频...", file=sys.stderr)
        async with httpx.AsyncClient(timeout=300, follow_redirects=True) as c:
            r = await c.get(video_url)
            r.raise_for_status()
            local_video.write_bytes(r.content)

        # Metadata
        print("[2/4] 提取元数据...", file=sys.stderr)
        meta = _probe_video(ffprobe_bin, local_video)
        dur = meta["duration"]
        if dur <= 0:
            print(json.dumps({"error": "无法获取视频时长"}))
            return

        # Segments
        print(f"[3/4] 构建分镜 ({dur:.1f}s)...", file=sys.stderr)
        segments = _build_segments(dur)

        # Frames
        print("[4/4] 提取关键帧...", file=sys.stderr)
        _extract_frames(ffmpeg_bin, local_video, segments)

        result = {
            "task_id": task_id,
            "duration": round(dur, 2),
            "resolution": f"{meta.get('width')}x{meta.get('height')}",
            "segment_count": len(segments),
            "segments": [
                {
                    "index": s["index"],
                    "start": round(s["start"], 2),
                    "end": round(s["end"], 2),
                    "frame_paths": s.get("frame_paths", []),
                }
                for s in segments
            ],
        }
        print(json.dumps(result, ensure_ascii=False, indent=2))
    finally:
        shutil.rmtree(temp_dir, ignore_errors=True)


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python process_video.py <video_url>", file=sys.stderr)
        sys.exit(1)
    asyncio.run(main(sys.argv[1]))

"""
上传本地视频文件到火山引擎 TOS 对象存储，返回签名 URL。

Usage:
    python scripts/video_upload.py "<file_path>" [bucket_name]

Examples:
    python scripts/video_upload.py "/path/to/video.mp4"
    python scripts/video_upload.py "/path/to/video.mp4" "my-bucket"
"""

import json
import os
import sys
from datetime import datetime

import tos
from tos import HttpMethodType

DEFAULT_BUCKET = "video-breakdown-uploads"
DEFAULT_REGION = "cn-beijing"


def video_upload_to_tos(file_path: str, bucket_name: str = None) -> dict:
    """
    将本地视频文件上传到 TOS，返回签名 URL。

    Args:
        file_path: 本地视频文件路径
        bucket_name: TOS 存储桶名称（可选）

    Returns:
        dict: 包含 video_url 或 error
    """
    if bucket_name is None:
        bucket_name = os.getenv("DATABASE_TOS_BUCKET") or os.getenv(
            "TOS_BUCKET", DEFAULT_BUCKET
        )
    region = os.getenv("DATABASE_TOS_REGION") or os.getenv("TOS_REGION", DEFAULT_REGION)

    # 检查文件
    if not os.path.exists(file_path):
        return {"error": f"文件不存在: {file_path}"}
    if not os.path.isfile(file_path):
        return {"error": f"路径不是文件: {file_path}"}

    file_size = os.path.getsize(file_path)
    max_size = 2 * 1024 * 1024 * 1024  # 2GB
    if file_size > max_size:
        return {"error": f"文件过大（{file_size / 1024 / 1024:.0f}MB），最大支持 2GB"}

    # 获取凭证
    access_key = os.getenv("VOLCENGINE_ACCESS_KEY", "")
    secret_key = os.getenv("VOLCENGINE_SECRET_KEY", "")
    session_token = ""

    if not access_key or not secret_key:
        try:
            from veadk.auth.veauth.utils import get_credential_from_vefaas_iam

            cred = get_credential_from_vefaas_iam()
            access_key = cred.access_key_id
            secret_key = cred.secret_access_key
            session_token = cred.session_token
        except Exception:
            pass

    if not access_key or not secret_key:
        return {
            "error": "缺少 TOS 访问凭证，请设置 VOLCENGINE_ACCESS_KEY 和 VOLCENGINE_SECRET_KEY"
        }

    # 自动生成 object_key
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = os.path.basename(file_path)
    object_key = f"video_breakdown/upload/{timestamp}_{filename}"

    # 上传
    client = None
    try:
        endpoint = f"tos-{region}.volces.com"
        client = tos.TosClientV2(
            ak=access_key,
            sk=secret_key,
            security_token=session_token,
            endpoint=endpoint,
            region=region,
        )

        # 检查桶
        try:
            client.head_bucket(bucket_name)
        except tos.exceptions.TosServerError as e:
            if e.status_code == 404:
                return {"error": f"TOS 存储桶 {bucket_name} 不存在"}
            raise

        print(f"上传中: {file_path} -> {bucket_name}/{object_key}", file=sys.stderr)
        client.put_object_from_file(
            bucket=bucket_name, key=object_key, file_path=file_path
        )

        # 生成签名 URL（7天有效）
        signed_url_output = client.pre_signed_url(
            http_method=HttpMethodType.Http_Method_Get,
            bucket=bucket_name,
            key=object_key,
            expires=604800,
        )

        return {
            "video_url": signed_url_output.signed_url,
            "bucket": bucket_name,
            "object_key": object_key,
            "file_size_mb": round(file_size / 1024 / 1024, 2),
            "message": "上传成功！使用 video_url 调用 process_video.py 进行视频分镜分析",
        }

    except tos.exceptions.TosClientError as e:
        return {"error": f"TOS 客户端错误: {e}"}
    except tos.exceptions.TosServerError as e:
        return {"error": f"TOS 服务端错误: {e.message}"}
    except Exception as e:
        return {"error": f"上传失败: {str(e)}"}
    finally:
        if client:
            client.close()


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python video_upload.py <file_path> [bucket_name]")
        sys.exit(1)

    path = sys.argv[1]
    bucket = sys.argv[2] if len(sys.argv) > 2 else None
    result = video_upload_to_tos(path, bucket)
    print(json.dumps(result, ensure_ascii=False, indent=2))