"""ASRProcessor - Audio-to-text conversion

Unified processor for B站官方字幕下载 and Whisper.cpp 语音识别.
Follows the pipeline pattern established in Phase 1/2.

Strategy:
1. Try official subtitle first (highest quality, no compute cost)
2. Fall back to Whisper.cpp if no subtitle available
3. Auto-enable speaker diarization if HF_TOKEN is configured
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import subprocess
from pathlib import Path
from typing import Optional, Dict, Any, List

import aiohttp

from kb_workflow.pipeline.config import get_settings
from kb_workflow.pipeline.schemas import ASRInput, ASROutput
from kb_workflow.pipeline.state import WorkflowStateManager

logger = logging.getLogger(__name__)


class WhisperNotFoundError(RuntimeError):
    """Whisper.cpp 路径未找到"""
    pass


def get_whisper_cpp_path() -> str:
    """获取 whisper.cpp 安装路径

    检查环境变量 WHISPER_CPP_PATH 或默认路径 ~/whisper.cpp

    Returns:
        whisper.cpp 目录路径

    Raises:
        WhisperNotFoundError: 如果路径不存在
    """
    # 检查环境变量
    path = os.environ.get('WHISPER_CPP_PATH')
    if path and os.path.exists(path):
        return os.path.expanduser(path)

    # 检查默认路径
    default = os.path.expanduser('~/whisper.cpp')
    if os.path.exists(default):
        return default

    raise WhisperNotFoundError(
        "未找到 whisper.cpp。请设置 WHISPER_CPP_PATH 环境变量或安装到 ~/whisper.cpp"
    )


def convert_to_wav(input_path: str, output_path: Optional[str] = None) -> str:
    """使用 ffmpeg 将音频/视频转换为 16kHz WAV 格式

    Args:
        input_path: 输入文件路径
        output_path: 输出文件路径（可选，默认为 input_path_audio.wav）

    Returns:
        输出 WAV 文件路径
    """
    if output_path is None:
        base = os.path.splitext(input_path)[0]
        output_path = f"{base}_audio.wav"

    cmd = [
        'ffmpeg', '-y', '-i', input_path,
        '-ar', '16000',  # 16kHz 采样率
        '-ac', '1',      # 单声道
        '-c:a', 'pcm_s16le',  # 16-bit PCM
        output_path
    ]

    subprocess.run(cmd, capture_output=True, check=True)
    return output_path


class ASRProcessor:
    """ASR 处理器 - 统一字幕下载和语音识别

    处理流程：
    1. 如果是 B站视频，尝试下载官方字幕
    2. 无官方字幕则调用 Whisper.cpp
    3. HF_TOKEN 配置时自动启用说话者分离
    4. 记录处理状态到 workflow_state.json

    Attributes:
        STEP_NAME: 步骤名称，用于状态记录
        settings: KB 配置
        state_manager: 状态管理器（可选）
    """

    STEP_NAME = "asr"

    def __init__(self, state_manager: Optional[WorkflowStateManager] = None):
        """初始化 ASR 处理器

        Args:
            state_manager: 状态管理器（可选，用于记录处理状态）
        """
        self.settings = get_settings()
        self.state_manager = state_manager

    async def process(self, input_data: ASRInput) -> ASROutput:
        """执行 ASR 处理

        策略：
        0. 先检查是否已有字幕文件（幂等性）
        1. 先尝试官方字幕（仅 B站）
        2. 无官方字幕则调用 Whisper
        3. 记录状态

        Args:
            input_data: ASR 输入数据

        Returns:
            ASR 处理结果

        Raises:
            Exception: 处理失败时抛出异常
        """
        result = None

        try:
            # Step 0: 检查是否已有字幕文件（幂等性）
            existing_result = self._check_existing_subtitle_files(input_data)
            if existing_result:
                logger.info(f"Found existing subtitle files, skipping processing")
                result = existing_result
            else:
                # Step 1: 尝试官方字幕（仅 B站）
                if input_data.platform == "bilibili":
                    result = await self._try_official_subtitle(input_data)

                # Step 2: 无官方字幕则调用 Whisper
                if result is None:
                    result = await self._run_whisper(input_data)

            # Step 3: 记录成功状态
            if self.state_manager:
                state = self.state_manager.load()
                if state:
                    summary = f"tool={result.tool}, segments={len(result.segments)}"
                    if result.speakers:
                        summary += f", speakers={len(result.speakers)}"
                    state.mark_step_completed(self.STEP_NAME, summary)
                    self.state_manager.save(state)

            return result

        except Exception as e:
            # 记录失败状态（只在所有重试都失败后）
            if self.state_manager:
                state = self.state_manager.load()
                if state:
                    state.mark_step_failed(self.STEP_NAME, str(e))
                    self.state_manager.save(state)
            raise

    def _check_existing_subtitle_files(self, input_data: ASRInput) -> Optional[ASROutput]:
        """检查是否已存在字幕文件（幂等性检查）

        Args:
            input_data: ASR 输入数据

        Returns:
            如果找到现有字幕文件返回 ASROutput，否则返回 None
        """
        output_dir = Path(input_data.content_dir)

        # 尝试两种命名格式：新格式（带 _subtitle 后缀）和旧格式
        base_names = [
            f"{input_data.platform}_{input_data.content_id}_subtitle",
            f"{input_data.platform}_{input_data.content_id}",
        ]

        for base_name in base_names:
            json_path = output_dir / f"{base_name}.json"
            srt_path = output_dir / f"{base_name}.srt"
            txt_path = output_dir / f"{base_name}.txt"

            if json_path.exists():
                # 读取现有字幕数据
                segments = []
                text = ""
                duration = 0.0

                try:
                    with open(json_path, 'r', encoding='utf-8') as f:
                        data = json.load(f)

                    # 解析字幕格式
                    if 'transcription' in data:
                        # Whisper 格式
                        for i, item in enumerate(data['transcription']):
                            offsets = item.get('offsets', {})
                            if isinstance(offsets, dict):
                                # 新格式：{"from": 0, "to": 5320}
                                start = offsets.get('from', 0) / 1000.0
                                end = offsets.get('to', 0) / 1000.0
                            else:
                                # 旧格式：[0, 5320]
                                start = offsets[0] / 1000.0 if len(offsets) > 0 else 0
                                end = offsets[1] / 1000.0 if len(offsets) > 1 else 0
                            segments.append({
                                'id': i,
                                'start': start,
                                'end': end,
                                'text': item.get('text', '')
                            })
                    elif 'body' in data:
                        # B站官方字幕格式
                        for i, item in enumerate(data['body']):
                            segments.append({
                                'id': i,
                                'start': item.get('from', 0),
                                'end': item.get('to', 0),
                                'text': item.get('content', '')
                            })

                    text = ' '.join(s['text'] for s in segments)
                    if segments:
                        duration = segments[-1]['end']

                    output_files = {'json': str(json_path)}
                    if srt_path.exists():
                        output_files['srt'] = str(srt_path)
                    if txt_path.exists():
                        output_files['txt'] = str(txt_path)

                    return ASROutput(
                        success=True,
                        tool="existing",
                        has_subtitle=True,
                        output_files=output_files,
                        segments=segments,
                        text=text,
                        duration_seconds=duration
                    )
                except Exception as e:
                    logger.warning(f"Failed to read existing subtitle file: {e}")
                    continue

        return None

    async def _try_official_subtitle(self, input_data: ASRInput) -> Optional[ASROutput]:
        """尝试下载 B站官方字幕

        Args:
            input_data: ASR 输入数据

        Returns:
            如果找到字幕返回 ASROutput，否则返回 None

        Note:
            无字幕返回 None（不是抛异常），让流程继续到 Whisper
        """
        # 使用 content_dir 作为输出目录（而不是从 video_path 推断）
        output_dir = Path(input_data.content_dir)
        # 使用 content_id 作为基础文件名
        base_name = f"{input_data.platform}_{input_data.content_id}_subtitle"

        # 输出文件路径
        txt_path = output_dir / f"{base_name}.txt"
        srt_path = output_dir / f"{base_name}.srt"
        json_path = output_dir / f"{base_name}.json"

        async with aiohttp.ClientSession() as session:
            try:
                # Step 1: 获取 CID
                bvid = input_data.content_id
                pagelist_url = f"https://api.bilibili.com/x/player/pagelist?bvid={bvid}"
                async with session.get(pagelist_url) as resp:
                    pagelist_data = await resp.json()

                pages = pagelist_data.get('data', [])
                if not pages:
                    logger.info(f"No pagelist data for {bvid}")
                    return None

                cid = pages[0].get('cid')
                if not cid:
                    logger.info(f"No cid found for {bvid}")
                    return None

                # Step 2: 获取字幕列表
                player_url = f"https://api.bilibili.com/x/player/v2?bvid={bvid}&cid={cid}"
                async with session.get(player_url) as resp:
                    player_data = await resp.json()

                subtitle_info = player_data.get('data', {}).get('subtitle', {})
                subtitles = subtitle_info.get('list', [])

                if not subtitles:
                    logger.info(f"No official subtitle for {bvid}")
                    return None

                # Step 3: 下载第一个字幕（优先 AI 生成的中文字幕）
                subtitle_url = None
                for sub in subtitles:
                    if sub.get('lan', '').startswith('zh'):
                        subtitle_url = sub.get('subtitle_url')
                        break

                if not subtitle_url and subtitles:
                    subtitle_url = subtitles[0].get('subtitle_url')

                if not subtitle_url:
                    return None

                # 补全 URL
                if subtitle_url.startswith('//'):
                    subtitle_url = f"https:{subtitle_url}"

                async with session.get(subtitle_url) as resp:
                    subtitle_content = await resp.json()

                # Step 4: 转换为输出格式
                segments = []
                text_parts = []
                body = subtitle_content.get('body', [])

                for i, item in enumerate(body):
                    start = item.get('from', 0)
                    end = item.get('to', 0)
                    content = item.get('content', '')

                    segments.append({
                        'id': i,
                        'start': start,
                        'end': end,
                        'text': content
                    })
                    text_parts.append(content)

                full_text = ' '.join(text_parts)
                duration = segments[-1]['end'] if segments else 0

                # Step 5: 写入文件
                # TXT 格式
                txt_path.write_text(full_text, encoding='utf-8')

                # SRT 格式
                srt_content = self._convert_to_srt(segments)
                srt_path.write_text(srt_content, encoding='utf-8')

                # JSON 格式
                json_data = {
                    'success': True,
                    'tool': 'subtitle',
                    'segments': segments,
                    'text': full_text,
                    'duration_seconds': duration
                }
                json_path.write_text(json.dumps(json_data, ensure_ascii=False, indent=2), encoding='utf-8')

                return ASROutput(
                    success=True,
                    tool="subtitle",
                    has_subtitle=True,
                    output_files={
                        'txt': str(txt_path),
                        'srt': str(srt_path),
                        'json': str(json_path)
                    },
                    segments=segments,
                    text=full_text,
                    duration_seconds=duration
                )

            except aiohttp.ClientError as e:
                logger.warning(f"Failed to fetch subtitle: {e}")
                return None

    def _convert_to_srt(self, segments: List[Dict]) -> str:
        """将字幕分段转换为 SRT 格式

        Args:
            segments: 字幕分段列表，每个包含 start, end, text

        Returns:
            SRT 格式字符串
        """
        lines = []
        for i, seg in enumerate(segments, 1):
            start_time = self._seconds_to_srt_time(seg['start'])
            end_time = self._seconds_to_srt_time(seg['end'])
            lines.append(f"{i}")
            lines.append(f"{start_time} --> {end_time}")
            lines.append(seg['text'])
            lines.append("")  # 空行

        return '\n'.join(lines)

    def _seconds_to_srt_time(self, seconds: float) -> str:
        """将秒数转换为 SRT 时间格式 (HH:MM:SS,mmm)"""
        hours = int(seconds // 3600)
        minutes = int((seconds % 3600) // 60)
        secs = int(seconds % 60)
        millis = int((seconds % 1) * 1000)
        return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"

    async def _run_whisper(self, input_data: ASRInput) -> ASROutput:
        """执行 Whisper 语音识别

        根据 HF_TOKEN 配置自动选择普通模式或说话者分离模式。

        Args:
            input_data: ASR 输入数据

        Returns:
            ASR 处理结果
        """
        hf_token = self.settings.hf_token.get_secret_value()
        if hf_token:
            return await self._run_whisper_with_diarization(input_data)
        else:
            return await self._run_whisper_basic(input_data)

    async def _run_whisper_basic(self, input_data: ASRInput) -> ASROutput:
        """执行普通 Whisper.cpp 语音识别

        Args:
            input_data: ASR 输入数据

        Returns:
            ASR 处理结果

        Raises:
            RuntimeError: Whisper.cpp 路径不存在或执行失败，或 video_path 为空
        """
        # 检查 video_path 是否有效
        if not input_data.video_path:
            raise RuntimeError(
                "video_path 为空，无法执行 Whisper 语音识别。"
                "请确保视频文件已下载（对于没有官方字幕的视频，需要先下载视频文件）。"
            )

        video_path = Path(input_data.video_path)
        output_dir = Path(input_data.content_dir)
        base_name = f"{input_data.platform}_{input_data.content_id}_whisper"

        # 获取 Whisper.cpp 路径
        try:
            whisper_path = get_whisper_cpp_path()
        except WhisperNotFoundError as e:
            raise RuntimeError(str(e))

        # 模型路径
        model_path = os.path.join(whisper_path, 'models', 'ggml-base.bin')
        if not os.path.exists(model_path):
            raise RuntimeError(f"Whisper model not found: {model_path}")

        # 转换音频格式
        audio_path = convert_to_wav(str(video_path))
        temp_audio_created = audio_path != str(video_path)

        try:
            # 构建 Whisper 命令
            output_base = os.path.join(output_dir, base_name)
            cmd = [
                os.path.join(whisper_path, 'whisper-cli'),
                '-m', model_path,
                '-f', audio_path,
                '-of', output_base,
                '-l', 'zh',
                '--output-txt', '--output-srt', '--output-json',
            ]

            # 执行 Whisper
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)

            if result.returncode != 0:
                raise RuntimeError(f"Whisper failed: {result.stderr}")

            # 解析输出
            return self._parse_whisper_output(output_dir, base_name)

        finally:
            # 清理临时音频文件
            if temp_audio_created and os.path.exists(audio_path):
                try:
                    os.remove(audio_path)
                except Exception:
                    pass  # 静默失败，不阻塞主流程

    def _parse_whisper_output(self, output_dir: Path, base_name: str) -> ASROutput:
        """解析 Whisper 输出文件

        Args:
            output_dir: 输出目录
            base_name: 文件基础名

        Returns:
            ASR 处理结果
        """
        txt_path = output_dir / f"{base_name}.txt"
        srt_path = output_dir / f"{base_name}.srt"
        json_path = output_dir / f"{base_name}.json"

        # 读取 JSON 输出
        segments = []
        text = ""
        duration = 0.0

        if json_path.exists():
            with open(json_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                # Whisper.cpp JSON 格式
                if 'transcription' in data:
                    for i, item in enumerate(data['transcription']):
                        offsets = item.get('offsets', {})
                        if isinstance(offsets, dict):
                            # 新格式：{"from": 0, "to": 5320}
                            start = offsets.get('from', 0) / 1000.0
                            end = offsets.get('to', 0) / 1000.0
                        else:
                            # 旧格式：[0, 5320]
                            start = offsets[0] / 1000.0 if len(offsets) > 0 else 0
                            end = offsets[1] / 1000.0 if len(offsets) > 1 else 0
                        segments.append({
                            'id': i,
                            'start': start,
                            'end': end,
                            'text': item.get('text', '')
                        })
                text = ' '.join(s['text'] for s in segments)
                if segments:
                    duration = segments[-1]['end']

        # 读取文本
        if txt_path.exists():
            text = txt_path.read_text(encoding='utf-8').strip()

        return ASROutput(
            success=True,
            tool="whisper",
            has_subtitle=False,
            output_files={
                'txt': str(txt_path),
                'srt': str(srt_path),
                'json': str(json_path)
            },
            segments=segments,
            text=text,
            duration_seconds=duration
        )

    async def _run_whisper_with_diarization(self, input_data: ASRInput) -> ASROutput:
        """执行带说话者分离的 Whisper

        Args:
            input_data: ASR 输入数据

        Returns:
            ASR 处理结果，包含说话者信息
        """
        # 先运行普通 Whisper
        basic_result = await self._run_whisper_basic(input_data)

        try:
            # 运行说话者分离
            diarization_result = run_diarization(input_data.video_path)

            # 合并结果
            merged_segments = assign_speakers_to_segments(
                basic_result.segments,
                diarization_result
            )

            # 提取说话者列表
            speakers = list(set(
                seg.get('speaker') for seg in merged_segments if seg.get('speaker')
            ))

            return ASROutput(
                success=True,
                tool="whisper-diarization",
                has_subtitle=False,
                output_files=basic_result.output_files,
                segments=merged_segments,
                text=basic_result.text,
                duration_seconds=basic_result.duration_seconds,
                speakers=sorted(speakers)
            )

        except ImportError:
            logger.warning("pyannote.audio not available, falling back to basic whisper")
            return basic_result
        except Exception as e:
            logger.warning(f"Diarization failed: {e}, falling back to basic whisper")
            return basic_result


def run_diarization(audio_path: str, num_speakers: Optional[int] = None) -> Dict[str, List[Dict]]:
    """运行 pyannote.audio 说话者分离

    Args:
        audio_path: 音频文件路径
        num_speakers: 预期说话者数量（可选）

    Returns:
        说话者时间线字典 {speaker_name: [{start, end}, ...]}

    Raises:
        ImportError: 如果 pyannote.audio 未安装
        RuntimeError: 如果 HF_TOKEN 未配置
    """
    hf_token = os.environ.get('HF_TOKEN')
    if not hf_token:
        raise RuntimeError("需要设置 HF_TOKEN 环境变量以使用说话者分离功能")

    from huggingface_hub import login
    from pyannote.audio import Pipeline
    import torch

    login(token=hf_token)

    pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    pipeline.to(device)

    diarization = pipeline(audio_path, num_speakers=num_speakers)

    # 解析结果
    speakers = {}
    for turn, _, speaker in diarization.itertracks(yield_label=True):
        if speaker not in speakers:
            speakers[speaker] = []
        speakers[speaker].append({
            "start": round(turn.start, 2),
            "end": round(turn.end, 2)
        })

    return speakers


def assign_speakers_to_segments(
    segments: List[Dict],
    diarization: Dict[str, List[Dict]]
) -> List[Dict]:
    """为每个段落分配说话者

    根据时间戳匹配说话者分离结果到 Whisper 段落。

    Args:
        segments: Whisper 段落列表
        diarization: 说话者分离结果

    Returns:
        带有 speaker 字段的段落列表
    """
    result = []

    for seg in segments:
        seg_start = seg.get('start', 0)
        seg_end = seg.get('end', 0)
        seg_mid = (seg_start + seg_end) / 2

        # 找到包含中点时间的说话者
        assigned_speaker = None
        for speaker, turns in diarization.items():
            for turn in turns:
                if turn['start'] <= seg_mid <= turn['end']:
                    assigned_speaker = speaker
                    break
            if assigned_speaker:
                break

        new_seg = dict(seg)
        new_seg['speaker'] = assigned_speaker
        result.append(new_seg)

    return result