"""
AutoVideo - 视频处理分析一键流程
方案: MiMo 2.5 (分镜) + Qwen3-ASR-Flash FileTrans (转录+字级时间戳)
"""
import subprocess
import json
import os
import sys
import io
import time
import base64
import requests

sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')


def phase2_preprocess(video_path, output_dir):
    """压缩视频 + 提取音频 + 提取关键帧"""
    name = os.path.splitext(os.path.basename(video_path))[0]
    compressed = os.path.join(output_dir, f"{name}_compressed.mp4")
    orig = os.path.getsize(video_path)
    subprocess.run([
        "ffmpeg", "-y", "-i", video_path,
        "-vf", "scale=-2:480", "-r", "15",
        "-c:v", "libx264", "-preset", "fast", "-crf", "28",
        "-c:a", "aac", "-b:a", "48k", "-ar", "16000", "-ac", "1",
        "-movflags", "+faststart", compressed
    ], capture_output=True, check=True)
    comp = os.path.getsize(compressed)
    print(f"  压缩: {orig/1024/1024:.1f}MB → {comp/1024/1024:.1f}MB ({comp/orig*100:.0f}%)")

    audio = os.path.join(output_dir, f"{name}_audio.wav")
    subprocess.run([
        "ffmpeg", "-y", "-i", video_path,
        "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1", audio
    ], capture_output=True, check=True)
    print(f"  音频: {os.path.getsize(audio)/1024:.0f} KB")

    frames_dir = os.path.join(output_dir, f"{name}_frames")
    os.makedirs(frames_dir, exist_ok=True)
    subprocess.run([
        "ffmpeg", "-y", "-i", video_path,
        "-vf", "fps=2,scale=480:-2", "-q:v", "5",
        os.path.join(frames_dir, "frame_%04d.jpg")
    ], capture_output=True, check=True)
    frames = len([f for f in os.listdir(frames_dir) if f.endswith('.jpg')])
    print(f"  关键帧: {frames} 帧")
    return compressed, audio


def phase3_mimo_scene(compressed_path):
    """MiMo 2.5 分镜分析 (OpenAI 格式)"""
    import openai
    with open(compressed_path, "rb") as f:
        video_b64 = base64.b64encode(f.read()).decode()
    client = openai.OpenAI(base_url="http://127.0.0.1:15721/v1", api_key="dummy-key")
    response = client.chat.completions.create(
        model="mimo-v2.5",
        messages=[{"role": "user", "content": [
            {"type": "video_url", "video_url": {"url": f"data:video/mp4;base64,{video_b64}"}},
            {"type": "text", "text": """请逐镜头分析这个视频，输出JSON：
{"video_summary":"概述","scenes":[{"scene_id":1,"start_time":"MM:SS.mmm","end_time":"MM:SS.mmm","duration_sec":数字,"scene_type":"类型","visual_description":"详细视觉描述","camera":"运镜","transition":"转场"}],"timeline":"时间线叙事总结"}"""}
        ]}],
        max_tokens=8192, temperature=0.1)
    text = response.choices[0].message.content
    clean = text.strip()
    if clean.startswith("```"): clean = clean.split("\n", 1)[1]
    if clean.endswith("```"): clean = clean.rsplit("```", 1)[0]
    try:
        return json.loads(clean.strip())
    except json.JSONDecodeError:
        return {"raw_text": text}


def phase4_qwen3_asr(audio_path):
    """Qwen3-ASR-Flash FileTrans: 转录 + 字级时间戳 (一步到位)"""
    import dashscope
    API_KEY = dashscope.api_key

    # 1. 上传获取临时URL
    upload = dashscope.Files.upload(file_path=audio_path, purpose="file-extract")
    file_id = upload.output["uploaded_files"][0]["file_id"]
    url = dashscope.Files.get(file_id=file_id).output["url"]

    # 2. 提交转录任务
    resp = requests.post(
        "https://dashscope.aliyuncs.com/api/v1/services/audio/asr/transcription",
        headers={"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json",
                 "X-DashScope-Async": "enable"},
        json={"model": "qwen3-asr-flash-filetrans",
              "input": {"file_url": url},
              "parameters": {"channel_id": [0], "enable_words": True, "enable_itn": False, "language": "zh"}}
    )
    task_id = resp.json()["output"]["task_id"]

    # 3. 轮询结果
    for i in range(30):
        time.sleep(2)
        result = requests.get(f"https://dashscope.aliyuncs.com/api/v1/tasks/{task_id}",
                              headers={"Authorization": f"Bearer {API_KEY}"}).json()
        if result["output"]["task_status"] == "SUCCEEDED":
            break
        elif result["output"]["task_status"] == "FAILED":
            raise Exception(f"转录失败: {result['output']}")

    # 4. 获取转录JSON
    tr_url = result["output"]["result"]["transcription_url"]
    return requests.get(tr_url).json()


def generate_srt(trans_data, output_path):
    """生成 SRT 字幕"""
    lines, idx = [], 1
    for ch in trans_data.get("transcripts", []):
        for sent in ch.get("sentences", []):
            b, e, t = sent["begin_time"], sent["end_time"], sent["text"]
            bh, bm, bs, bms = b//3600000, (b%3600000)//60000, (b%60000)//1000, b%1000
            eh, em, es, ems = e//3600000, (e%3600000)//60000, (e%60000)//1000, e%1000
            lines.extend([f"{idx}",
                f"{bh:02d}:{bm:02d}:{bs:02d},{bms:03d} --> {eh:02d}:{em:02d}:{es:02d},{ems:03d}",
                t, ""])
            idx += 1
    with open(output_path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))


def print_timestamps(trans_data):
    """打印字级时间戳"""
    for ch in trans_data.get("transcripts", []):
        total = 0
        for sent in ch.get("sentences", []):
            print(f"\n  [{sent['begin_time']:>6}ms - {sent['end_time']:>6}ms] {sent['text']}")
            for w in sent.get("words", []):
                print(f"    [{w['begin_time']:>6}ms - {w['end_time']:>6}ms] {w['text']}{w.get('punctuation','')}")
                total += 1
        print(f"\n  总字数: {total}")


def run_pipeline(video_path, output_dir=None):
    name = os.path.splitext(os.path.basename(video_path))[0]
    if output_dir is None:
        output_dir = os.path.join(os.path.dirname(video_path), f"{name}_analysis")
    os.makedirs(output_dir, exist_ok=True)

    print(f"{'='*60}\nAutoVideo: {name}\n{'='*60}")

    # Phase 2
    print(f"\n▶ Phase 2: 压缩预处理")
    compressed, audio = phase2_preprocess(video_path, output_dir)

    # Phase 3
    print(f"\n▶ Phase 3: MiMo 2.5 分镜分析")
    try:
        scene = phase3_mimo_scene(compressed)
        with open(os.path.join(output_dir, "scene.json"), "w", encoding="utf-8") as f:
            json.dump(scene, f, ensure_ascii=False, indent=2)
        for s in scene.get("scenes", []):
            print(f"  [{s.get('start_time','?')}→{s.get('end_time','?')}] {s.get('scene_type','')}")
    except Exception as e:
        print(f"  错误: {e}")

    # Phase 4 (转录 + 字级时间戳 一步到位)
    print(f"\n▶ Phase 4: Qwen3-ASR 转录 + 字级时间戳")
    try:
        ts = phase4_qwen3_asr(audio)
        with open(os.path.join(output_dir, "timestamps.json"), "w", encoding="utf-8") as f:
            json.dump(ts, f, ensure_ascii=False, indent=2)
        generate_srt(ts, os.path.join(output_dir, "subtitles.srt"))
        print_timestamps(ts)
    except Exception as e:
        print(f"  错误: {e}")

    # 汇总
    print(f"\n{'='*60}\n完成! 输出: {output_dir}\n{'='*60}")
    for f in sorted(os.listdir(output_dir)):
        fp = os.path.join(output_dir, f)
        if os.path.isfile(fp):
            print(f"  {f:40s} {os.path.getsize(fp)/1024:>6.1f} KB")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python pipeline.py <video_path> [output_dir]")
        sys.exit(1)
    run_pipeline(sys.argv[1], sys.argv[2] if len(sys.argv) > 2 else None)