#!/usr/bin/env python3
"""
字幕生成器 - 为无字幕视频自动生成 SRT 格式字幕

使用 OpenAI Whisper 进行语音识别，支持本地离线处理
"""

import os
import time
import hashlib
import logging
from pathlib import Path
from typing import Optional, Callable, Dict, Any

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class AudioDownloader:
    """使用 yt-dlp 下载视频音频"""

    def __init__(self):
        """初始化音频下载器"""
        self.ydl_opts = {
            'format': 'bestaudio/best',
            'postprocessors': [{
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'mp3',
                'preferredquality': '192',
            }],
            'quiet': True,
            'no_warnings': True,
        }

    def download_audio(self, url: str, output_path: str) -> str:
        """
        下载视频音频

        Args:
            url: 视频URL
            output_path: 输出音频路径（可以是模板，如 'cache/audio/%(id)s.mp3'）

        Returns:
            str: 实际保存的音频文件路径

        Raises:
            Exception: 下载失败时抛出异常
        """
        try:
            import yt_dlp

            ydl_opts = self.ydl_opts.copy()
            ydl_opts['outtmpl'] = output_path

            logger.info(f"开始下载音频: {url}")
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                ydl.download([url])

            # 替换扩展名
            actual_path = output_path.replace('.%(ext)s', '.mp3')
            if not os.path.exists(actual_path):
                # 尝试其他可能的文件名
                actual_path = output_path.rsplit('.', 1)[0] + '.mp3'

            if os.path.exists(actual_path):
                logger.info(f"音频下载成功: {actual_path}")
                return actual_path
            else:
                raise FileNotFoundError(f"下载的音频文件未找到: {actual_path}")

        except Exception as e:
            logger.error(f"音频下载失败: {e}")
            raise


class WhisperTranscriber:
    """Whisper 语音识别器"""

    # 模型选择
    MODELS = {
        'tiny': 39000000,      # 最快，准确度最低 (~39M)
        'base': 74000000,      # 推荐，速度/准确度平衡 (~74M)
        'small': 244000000,    # 准确度更高 (~244M)
        'medium': 769000000,   # 接近专业级 (~769M)
        'large': 1550000000,   # 最高准确度 (~1.55GB)
    }

    def __init__(self, model_size: str = 'base'):
        """
        初始化识别器

        Args:
            model_size: 模型大小，默认 'base'
        """
        self.model_size = model_size
        self.model = None
        self.device = None

    def load_model(self):
        """加载 Whisper 模型（懒加载）"""
        if self.model is None:
            logger.info(f"正在加载 Whisper 模型 ({self.model_size})...")
            start_time = time.time()

            try:
                import whisper
                import torch

                # 自动检测 GPU
                self.device = "cuda" if torch.cuda.is_available() else "cpu"
                logger.info(f"使用设备: {self.device}")

                self.model = whisper.load_model(
                    self.model_size,
                    device=self.device
                )

                load_time = time.time() - start_time
                logger.info(f"模型加载完成 - 耗时: {load_time:.1f}秒")

            except Exception as e:
                logger.error(f"模型加载失败: {e}")
                raise

    def transcribe(
        self,
        audio_path: str,
        language: Optional[str] = None,
        progress_callback: Optional[Callable[[str, float], None]] = None
    ) -> Dict[str, Any]:
        """
        语音识别

        Args:
            audio_path: 音频文件路径
            language: 语言代码（'zh'、'en'、None=自动检测）
            progress_callback: 进度回调函数 callback(message, progress_percent)

        Returns:
            dict: Whisper 识别结果，包含 'text' 和 'segments'

        Raises:
            Exception: 识别失败时抛出异常
        """
        self.load_model()

        logger.info(f"开始语音识别: {audio_path}")
        if language:
            logger.info(f"指定语言: {language}")
        else:
            logger.info("语言: 自动检测")

        start_time = time.time()

        try:
            result = self.model.transcribe(
                audio_path,
                language=language,
                task='transcribe',
                verbose=False,
                fp16=False  # CPU 上不使用 FP16
            )

            duration = time.time() - start_time
            logger.info(f"识别完成 - 耗时: {duration:.1f}秒")

            return result

        except Exception as e:
            logger.error(f"语音识别失败: {e}")
            raise

    def get_segments(self, result: Dict[str, Any]) -> list:
        """
        获取分段信息

        Args:
            result: Whisper 识别结果

        Returns:
            list: segments 列表
        """
        return result.get('segments', [])


class SubtitleFormatter:
    """字幕格式化器（SRT 格式）"""

    @staticmethod
    def to_srt(segments: list) -> str:
        """
        将 Whisper segments 转换为 SRT 格式

        Args:
            segments: Whisper segments 列表

        Returns:
            str: SRT 格式文本
        """
        srt_lines = []

        for i, segment in enumerate(segments, 1):
            # SRT 格式：
            # 序号
            # 开始时间 --> 结束时间
            # 文本
            # 空行

            start_time = SubtitleFormatter._seconds_to_srt_time(segment['start'])
            end_time = SubtitleFormatter._seconds_to_srt_time(segment['end'])
            text = segment['text'].strip()

            srt_lines.append(str(i))
            srt_lines.append(f"{start_time} --> {end_time}")
            srt_lines.append(text)
            srt_lines.append("")

        return "\n".join(srt_lines)

    @staticmethod
    def _seconds_to_srt_time(seconds: float) -> str:
        """
        将秒数转换为 SRT 时间格式 (HH:MM:SS,mmm)

        Args:
            seconds: 秒数

        Returns:
            str: SRT 时间格式
        """
        hours = int(seconds // 3600)
        minutes = int((seconds % 3600) // 60)
        secs = int(seconds % 60)
        millis = int((seconds - int(seconds)) * 1000)

        return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"


class SubtitleCache:
    """字幕缓存管理器"""

    def __init__(self, cache_dir: str = 'cache/subtitles'):
        """
        初始化缓存

        Args:
            cache_dir: 缓存目录
        """
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(parents=True, exist_ok=True)

    def get_cache_key(self, url: str) -> str:
        """
        生成缓存键

        Args:
            url: 视频URL

        Returns:
            str: 缓存键（SHA256 哈希的前8位）
        """
        return hashlib.sha256(url.encode()).hexdigest()[:8]

    def get_cache_path(self, url: str) -> str:
        """
        获取缓存文件路径

        Args:
            url: 视频URL

        Returns:
            str: 缓存文件路径
        """
        cache_key = self.get_cache_key(url)
        return str(self.cache_dir / f"{cache_key}.srt")

    def get(self, url: str) -> Optional[str]:
        """
        获取缓存字幕

        Args:
            url: 视频URL

        Returns:
            str | None: 缓存的字幕内容，如果不存在返回 None
        """
        cache_path = self.get_cache_path(url)

        if os.path.exists(cache_path):
            logger.info(f"使用缓存字幕: {cache_path}")
            try:
                with open(cache_path, 'r', encoding='utf-8') as f:
                    return f.read()
            except Exception as e:
                logger.warning(f"读取缓存失败: {e}")
                return None

        return None

    def set(self, url: str, subtitle: str) -> str:
        """
        保存缓存字幕

        Args:
            url: 视频URL
            subtitle: 字幕内容

        Returns:
            str: 缓存文件路径
        """
        cache_path = self.get_cache_path(url)

        try:
            with open(cache_path, 'w', encoding='utf-8') as f:
                f.write(subtitle)
            logger.info(f"字幕已缓存: {cache_path}")
            return cache_path
        except Exception as e:
            logger.error(f"保存缓存失败: {e}")
            raise

    def exists(self, url: str) -> bool:
        """检查缓存是否存在"""
        return os.path.exists(self.get_cache_path(url))

    def clear_cache(self):
        """清空所有缓存"""
        try:
            for file in self.cache_dir.glob('*.srt'):
                file.unlink()
            logger.info("缓存已清空")
        except Exception as e:
            logger.error(f"清空缓存失败: {e}")


class SubtitleGenerator:
    """字幕生成器统一接口"""

    def __init__(
        self,
        model_size: str = 'base',
        cache_dir: str = 'cache/subtitles',
        audio_cache_dir: str = 'cache/audio'
    ):
        """
        初始化生成器

        Args:
            model_size: Whisper 模型大小
            cache_dir: 字幕缓存目录
            audio_cache_dir: 音频缓存目录
        """
        self.audio_downloader = AudioDownloader()
        self.transcriber = WhisperTranscriber(model_size)
        self.formatter = SubtitleFormatter()
        self.cache = SubtitleCache(cache_dir)
        self.audio_cache_dir = Path(audio_cache_dir)
        self.audio_cache_dir.mkdir(parents=True, exist_ok=True)

    def generate(
        self,
        url: str,
        content_id: str = None,
        language: Optional[str] = None,
        use_cache: bool = True,
        keep_audio: bool = False
    ) -> str:
        """
        生成字幕

        Args:
            url: 视频URL
            content_id: 内容ID（用于音频文件命名）
            language: 语言代码（None=自动检测）
            use_cache: 是否使用缓存
            keep_audio: 是否保留音频文件

        Returns:
            str: SRT 格式字幕

        Raises:
            Exception: 生成失败时抛出异常
        """
        # 检查缓存
        if use_cache and self.cache.exists(url):
            logger.info(f"使用缓存字幕: {url}")
            cached_subtitle = self.cache.get(url)
            if cached_subtitle:
                return cached_subtitle

        logger.info(f"开始生成字幕: {url}")
        audio_path = None

        try:
            # 生成 content_id
            if content_id is None:
                content_id = self.cache.get_cache_key(url)

            # 下载音频
            audio_path = str(self.audio_cache_dir / f"{content_id}.mp3")
            logger.info(f"下载音频到: {audio_path}")
            audio_path = self.audio_downloader.download_audio(url, audio_path)

            # 语音识别
            result = self.transcriber.transcribe(audio_path, language=language)
            segments = self.transcriber.get_segments(result)

            if not segments:
                logger.warning("未识别到任何语音内容")
                return ""

            # 转换为 SRT
            srt_text = self.formatter.to_srt(segments)
            logger.info(f"SRT 字幕生成成功，共 {len(segments)} 段")

            # 保存缓存
            self.cache.set(url, srt_text)

            # 清理音频文件
            if not keep_audio and os.path.exists(audio_path):
                os.remove(audio_path)
                logger.info(f"已删除音频文件: {audio_path}")

            return srt_text

        except Exception as e:
            logger.error(f"字幕生成失败: {e}")
            # 清理失败的音频文件
            if audio_path and os.path.exists(audio_path) and not keep_audio:
                try:
                    os.remove(audio_path)
                except:
                    pass
            raise


def generate_subtitle(
    url: str,
    model_size: str = 'base',
    language: Optional[str] = None,
    cache_dir: str = 'cache/subtitles',
    audio_cache_dir: str = 'cache/audio'
) -> str:
    """
    便捷函数：生成字幕

    Args:
        url: 视频URL
        model_size: Whisper 模型大小
        language: 语言代码（None=自动检测）
        cache_dir: 字幕缓存目录
        audio_cache_dir: 音频缓存目录

    Returns:
        str: SRT 格式字幕
    """
    generator = SubtitleGenerator(
        model_size=model_size,
        cache_dir=cache_dir,
        audio_cache_dir=audio_cache_dir
    )
    return generator.generate(url, language=language)


if __name__ == '__main__':
    # 测试示例
    import sys

    if len(sys.argv) < 2:
        print("用法: python subtitle_generator.py <视频URL> [语言]")
        print("示例: python subtitle_generator.py https://www.bilibili.com/video/BV1xx411c7mD zh")
        sys.exit(1)

    url = sys.argv[1]
    language = sys.argv[2] if len(sys.argv) > 2 else None

    print("=" * 60)
    print("字幕生成器")
    print("=" * 60)

    try:
        subtitle = generate_subtitle(url, language=language)
        print("\n" + "=" * 60)
        print("生成成功！字幕预览（前500字符）：")
        print("=" * 60)
        print(subtitle[:500])
        if len(subtitle) > 500:
            print("\n...")
        print(f"\n总字符数: {len(subtitle)}")
    except Exception as e:
        print(f"\n错误: {e}")
        sys.exit(1)