#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
下载视频片段用于 OCR 提取硬字幕
策略：每隔 N 秒截取一帧，而不是下载整个视频
"""
import asyncio
import json
import os
import cv2
import numpy as np
from playwright.async_api import async_playwright
from urllib.request import Request, urlopen

async def extract_frames_for_ocr(bvid, output_dir="/tmp/bilibili_frames", interval_seconds=10, max_frames=200):
    """
    提取视频帧用于 OCR

    Args:
        bvid: B站视频 BV 号
        output_dir: 输出目录
        interval_seconds: 截帧间隔（秒）
        max_frames: 最大截取帧数

    Returns:
        str: 输出目录路径
    """
    url = f"https://www.bilibili.com/video/{bvid}"

    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)

    print(f"正在提取视频帧: {url}")
    print(f"间隔: {interval_seconds}秒，最大帧数: {max_frames}")
    print("=" * 80)

    async with async_playwright() as playwright:
        browser = await playwright.chromium.launch(headless=True)
        context = await browser.new_context()
        page = await context.new_page()

        # 获取视频信息
        print("\n1. 获取视频信息...")
        try:
            api_url = "https://api.bilibili.com/x/web-interface/view"
            response = await page.request.get(api_url, params={'bvid': bvid})
            data = await response.json()

            if data.get('code') != 0:
                print(f"   ❌ 失败: {data.get('message')}")
                await browser.close()
                return None

            video_data = data.get('data', {})
            title = video_data.get('title')
            aid = video_data.get('aid')
            cid = video_data.get('cid')
            duration = video_data.get('duration')
            print(f"   标题: {title}")
            print(f"   时长: {duration}秒")
        except Exception as e:
            print(f"   ❌ 异常: {e}")
            await browser.close()
            return None

        # 获取视频流 URL
        print("\n2. 获取视频流 URL...")
        try:
            playurl_api = "https://api.bilibili.com/x/player/playurl"
            params = {
                'bvid': bvid,
                'cid': cid,
                'qn': 80,  # 1080P
                'otype': 'json',
                'fnver': 0,
                'fnval': 0,  # FLV 格式
            }

            playurl_response = await page.request.get(playurl_api, params=params)

            if not playurl_response.ok or playurl_response.status != 200:
                print(f"   ❌ HTTP 错误: {playurl_response.status}")
                await browser.close()
                return None

            playurl_data = await playurl_response.json()

            if playurl_data.get('code') != 0:
                print(f"   ❌ API 错误: {playurl_data.get('message')}")
                await browser.close()
                return None

            play_data = playurl_data.get('data', {})
            durl = play_data.get('durl', [])

            if not durl:
                print(f"   ❌ 没有找到视频流")
                await browser.close()
                return None

            # 获取第一个分段的 URL
            first_segment_url = durl[0].get('url')
            print(f"   ✅ 获取到视频流 URL")

        except Exception as e:
            print(f"   ❌ 异常: {e}")
            import traceback
            traceback.print_exc()
            await browser.close()
            return None

        # 使用 OpenCV 直接从 URL 读取视频并提取帧
        print(f"\n3. 提取视频帧...")

        try:
            # 添加请求头
            headers = {
                'Referer': 'https://www.bilibili.com/',
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36',
            }

            # 使用 urllib 下载视频到临时文件（只下载前 50MB）
            print(f"   正在下载视频（前 50MB）...")

            req = Request(first_segment_url, headers=headers)

            # 下载视频数据（限制 50MB）
            chunk_size = 8192
            max_bytes = 50 * 1024 * 1024  # 50MB
            video_data = bytearray()

            with urlopen(req, timeout=60) as response:
                while True:
                    chunk = response.read(chunk_size)
                    if not chunk or len(video_data) >= max_bytes:
                        break
                    video_data.extend(chunk)

            print(f"   已下载 {len(video_data) / 1024 / 1024:.2f}MB")

            # 保存为临时文件
            temp_video_path = os.path.join(output_dir, "temp_video.flv")
            with open(temp_video_path, 'wb') as f:
                f.write(bytes(video_data))

            # 使用 OpenCV 打开视频
            cap = cv2.VideoCapture(temp_video_path)

            if not cap.isOpened():
                print(f"   ❌ 无法打开视频文件")
                await browser.close()
                return None

            # 获取视频帧率
            fps = cap.get(cv2.CAP_PROP_FPS)
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            print(f"   帧率: {fps:.2f} fps")
            print(f"   总帧数: {total_frames}")

            # 计算需要提取的帧
            interval_frames = int(fps * interval_seconds)
            frame_count = 0
            extracted_count = 0

            print(f"\n   开始提取帧...")

            while True:
                # 跳到目标帧
                cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)

                # 读取帧
                ret, frame = cap.read()

                if not ret:
                    break

                # 保存帧
                frame_path = os.path.join(output_dir, f"frame_{extracted_count:04d}_{frame_count}.jpg")
                cv2.imwrite(frame_path, frame)
                extracted_count += 1

                print(f"      提取第 {extracted_count} 帧 (时间: {frame_count / fps:.1f}s)")

                # 检查是否达到最大帧数
                if extracted_count >= max_frames:
                    print(f"      已达到最大帧数限制")
                    break

                # 跳到下一帧
                frame_count += interval_frames

            cap.release()

            # 删除临时视频文件
            os.remove(temp_video_path)

            print(f"\n4. 完成！")
            print(f"   成功提取 {extracted_count} 帧图像")
            print(f"   输出目录: {output_dir}")

            await browser.close()
            return output_dir

        except Exception as e:
            print(f"   ❌ 提取帧失败: {e}")
            import traceback
            traceback.print_exc()
            await browser.close()
            return None

if __name__ == '__main__':
    import sys

    bvid = "BV1KtFzzqEUG"

    result = asyncio.run(extract_frames_for_ocr(
        bvid,
        interval_seconds=10,  # 每 10 秒提取一帧
        max_frames=100  # 最多提取 100 帧
    ))

    if result:
        print(f"\n✅ 提取完成！")
        print(f"输出目录: {result}")
        print(f"\n下一步：使用 OCR 工具识别字幕")
    else:
        print(f"\n❌ 提取失败")
        sys.exit(1)