"""ASR Corrector Module

使用 LLM 纠正语音识别错误。复用 SimpleAIClient。

功能：
- 纠正明显的语音识别错误（同音字、漏字、错字）
- 分批处理长文本
- AI 不可用时返回原文
"""

from __future__ import annotations

import logging
from typing import Optional, Protocol

logger = logging.getLogger(__name__)


class AIClientProtocol(Protocol):
    """AI 客户端协议"""

    def generate(self, prompt: str, max_tokens: int = 16000) -> str:
        """生成文本"""
        ...


# 纠错 Prompt 模板
CORRECTION_PROMPT = """你是一个文本校对专家。以下是语音识别生成的字幕，可能包含识别错误。
请只纠正明显的错误（如同音字、漏字、错字），不要改变原意和风格。

原文：
{text}

纠正后的文本（只输出纠正后的文本，不要解释）："""


class ASRCorrector:
    """ASR 纠错器

    使用 LLM 纠正语音识别中的明显错误。

    Attributes:
        ai_client: AI 客户端（必须实现 generate 方法）
        batch_size: 分批处理的字符数，默认 5000
    """

    def __init__(
        self,
        ai_client: Optional[AIClientProtocol] = None,
        batch_size: int = 5000
    ):
        """初始化纠错器

        Args:
            ai_client: AI 客户端实例，如果为 None 则不进行纠错
            batch_size: 分批处理的字符数
        """
        self.ai_client = ai_client
        self.batch_size = batch_size

    def correct(self, text: str) -> str:
        """纠正文本中的语音识别错误

        Args:
            text: 原始文本

        Returns:
            纠正后的文本。如果 AI 不可用或文本为空，返回原文。

        Example:
            >>> corrector = ASRCorrector(ai_client=my_client)
            >>> result = corrector.correct("什么是机器学习")
            >>> print(result)
            什么是机器学习
        """
        # 空文本直接返回
        if not text or not text.strip():
            return text or ""

        # AI 不可用，返回原文
        if self.ai_client is None:
            logger.info("AI client not available, returning original text")
            return text

        # 短文本直接处理
        if len(text) <= self.batch_size:
            return self._correct_batch(text)

        # 长文本分批处理
        return self._correct_long_text(text)

    def _correct_batch(self, text: str) -> str:
        """处理单个批次的文本

        Args:
            text: 文本内容

        Returns:
            纠正后的文本
        """
        try:
            prompt = CORRECTION_PROMPT.format(text=text)
            # ASR 纠错需要直接结果，不需要 AI 的分析过程，因此禁用思考模式
            result = self.ai_client.generate(prompt, max_tokens=len(text) + 500, disable_reasoning=True)
            return result.strip()
        except Exception as e:
            logger.warning(f"ASR correction failed: {e}, returning original text")
            return text

    def _correct_long_text(self, text: str) -> str:
        """分批处理长文本

        按句子分割，每批不超过 batch_size 字符。

        Args:
            text: 长文本内容

        Returns:
            纠正后的文本
        """
        # 按句子分割（简单分割，基于句号、问号、感叹号）
        sentences = []
        current = ""
        for char in text:
            current += char
            if char in '。！？\n':
                sentences.append(current)
                current = ""

        # 添加最后一个不完整的句子
        if current:
            sentences.append(current)

        # 分批处理
        batches = []
        current_batch = []
        current_length = 0

        for sentence in sentences:
            if current_length + len(sentence) > self.batch_size:
                if current_batch:
                    batches.append(''.join(current_batch))
                current_batch = [sentence]
                current_length = len(sentence)
            else:
                current_batch.append(sentence)
                current_length += len(sentence)

        if current_batch:
            batches.append(''.join(current_batch))

        # 处理每个批次
        corrected_batches = []
        for batch in batches:
            corrected = self._correct_batch(batch)
            corrected_batches.append(corrected)

        return ''.join(corrected_batches)