"""语义分段引擎

使用 LLM 实现智能视频字幕语义分段。

功能：
- 视频类型识别（知识类/故事类/访谈类/新闻类/vlog 等）
- 语义边界检测
- 内容完整性保护
- 主题标题生成
"""

from __future__ import annotations

import json
import logging
import re
from typing import List, Optional, Protocol

from pydantic import BaseModel, Field, field_validator, ValidationError

logger = logging.getLogger(__name__)


# ============================================================================
# AI Client Protocol
# ============================================================================

class AIClientProtocol(Protocol):
    """AI 客户端协议（与 ASRCorrector 保持一致）"""

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


# ============================================================================
# Pydantic Models
# ============================================================================


def extract_json_from_llm_response(response: str) -> dict:
    """从 LLM 响应中提取 JSON

    处理以下情况：
    - 纯 JSON
    - 包裹在 ```json ``` 中的 JSON
    - 前后包含额外说明文本的 JSON

    Args:
        response: LLM 响应文本

    Returns:
        解析后的 JSON 字典

    Raises:
        ValueError: 无法提取有效 JSON
    """
    # 尝试直接解析
    try:
        return json.loads(response.strip())
    except json.JSONDecodeError:
        pass

    # 尝试提取 ```json ``` 代码块
    json_match = re.search(r'```json\s*(.*?)\s*```', response, re.DOTALL)
    if json_match:
        try:
            return json.loads(json_match.group(1))
        except json.JSONDecodeError:
            pass

    # 尝试提取任何 JSON 对象
    json_match = re.search(r'\{.*\}', response, re.DOTALL)
    if json_match:
        try:
            return json.loads(json_match.group())
        except json.JSONDecodeError as e:
            raise ValueError(f"无法从 LLM 响应中提取有效 JSON: {e}")

    raise ValueError("LLM 响应中未找到 JSON 内容")

class VideoSegment(BaseModel):
    """视频分段

    Attributes:
        title: 分段标题（可选）
        content: 分段内容（必填）
        start_index: 起始字幕索引
        end_index: 结束字幕索引
    """
    title: Optional[str] = Field(
        None,
        description="分段标题，简洁描述型（5-15字）"
    )
    content: str = Field(
        ...,
        description="分段内容，合并后的字幕文本"
    )
    start_index: int = Field(
        ...,
        ge=0,
        description="在原文中的起始字符位置"
    )
    end_index: int = Field(
        ...,
        ge=0,
        description="在原文中的结束字符位置"
    )

    @field_validator('title')
    @classmethod
    def validate_title_length(cls, v: Optional[str]) -> Optional[str]:
        """验证标题长度（5-15字）"""
        if v is not None and len(v) > 15:
            raise ValueError(f"标题长度不能超过15字，当前为{len(v)}字")
        return v

    model_config = {
        "json_schema_extra": {
            "examples": [
                {
                    "title": "机器学习概述",
                    "content": "机器学习是人工智能的一个分支...",
                    "start_index": 1,
                    "end_index": 10
                }
            ]
        }
    }


class SegmentationResult(BaseModel):
    """分段结果

    Attributes:
        video_type: 视频类型（知识类/故事类/访谈类/新闻类/vlog 等）
        segments: 分段列表
        actual_segment_count: 实际分段数量
    """
    video_type: str = Field(
        ...,
        description="识别出的视频类型"
    )
    segments: List[VideoSegment] = Field(
        default_factory=list,
        description="分段列表"
    )
    actual_segment_count: int = Field(
        ...,
        ge=0,
        description="实际分段数量"
    )

    model_config = {
        "json_schema_extra": {
            "examples": [
                {
                    "video_type": "知识类",
                    "segments": [
                        {
                            "title": "机器学习概述",
                            "content": "...",
                            "start_index": 1,
                            "end_index": 10
                        }
                    ],
                    "actual_segment_count": 5
                }
            ]
        }
    }


# ============================================================================
# Video Type Detection
# ============================================================================

# 视频类型关键词映射
VIDEO_TYPE_KEYWORDS = {
    "知识类": [
        "教程", "课程", "教学", "入门", "基础", "进阶", "实战",
        "编程", "开发", "技术", "学习", "教育", "培训", "讲座"
    ],
    "故事类": [
        "影视", "电影", "电视剧", "动漫", "动画", "小说", "故事",
        "剧情", "短片", "纪录片"
    ],
    "访谈类": [
        "访谈", "对话", "采访", "专访", "聊天", "脱口秀", "播客"
    ],
    "新闻类": [
        "新闻", "报道", "资讯", "时事", "热点", "快讯"
    ],
    "vlog": [
        "vlog", "日常", "生活", "旅行", "美食", "探店",
        "分享", "体验"
    ]
}

# 默认视频类型
DEFAULT_VIDEO_TYPE = "知识类"


# ============================================================================
# Semantic Segmentation Engine
# ============================================================================

class SemanticSegmentationEngine:
    """语义分段引擎

    使用 LLM 实现智能视频字幕语义分段。

    功能：
    - 识别视频类型（从元数据推断）
    - 构建包含类型信息的 Prompt
    - 调用 LLM 进行语义分段
    - 验证和返回分段结果

    Attributes:
        ai_client: AI 客户端（可选，暂未实现完整 LLM 调用）
    """

    def __init__(
        self,
        ai_client: Optional[AIClientProtocol] = None,
        fallback_segmenter: Optional['SemanticSegmenter'] = None
    ):
        """初始化语义分段引擎

        Args:
            ai_client: AI 客户端实例，如果为 None 则尝试创建 SimpleAIClient
            fallback_segmenter: 降级时使用的规则分段器
        """
        self.ai_client = ai_client
        self.fallback_segmenter = fallback_segmenter

        # Lazy AI Client 初始化（与 ASRCorrector 相同模式）
        if self.ai_client is None:
            try:
                from kb_workflow.deep_analyzer import SimpleAIClient
                self.ai_client = SimpleAIClient()
                if not self.ai_client.api_key:
                    logger.info("AI API key not configured")
                    self.ai_client = None
            except Exception as e:
                logger.warning(f"Failed to initialize AI client: {e}")
                self.ai_client = None

        # 初始化降级分段器
        if self.fallback_segmenter is None:
            from kb_workflow.srt.segmenter import SemanticSegmenter
            self.fallback_segmenter = SemanticSegmenter()

    def _infer_video_type(self, metadata: dict) -> str:
        """从元数据推断视频类型

        Args:
            metadata: 视频元数据，包含 title, category, tags 等字段

        Returns:
            推断出的视频类型，默认为"知识类"

        Examples:
            >>> engine = SemanticSegmentationEngine()
            >>> engine._infer_video_type({"title": "Python教程"})
            '知识类'
            >>> engine._infer_video_type({"category": "影视"})
            '故事类'
        """
        if not metadata:
            return DEFAULT_VIDEO_TYPE

        # 提取所有文本字段
        title = metadata.get("title", "") or ""
        category = metadata.get("category", "") or ""
        tags = metadata.get("tags", []) or []
        description = metadata.get("description", "") or ""

        # 合并所有文本用于匹配
        all_text = " ".join([title, category, description] + list(tags))

        # 如果没有有效文本，返回默认类型
        if not all_text or all_text.isspace():
            return DEFAULT_VIDEO_TYPE

        # 按优先级匹配类型关键词
        for video_type, keywords in VIDEO_TYPE_KEYWORDS.items():
            for keyword in keywords:
                if keyword in all_text:
                    logger.debug(
                        f"Detected video type '{video_type}' "
                        f"from keyword '{keyword}'"
                    )
                    return video_type

        # 未匹配到任何类型，返回默认
        logger.debug("No specific video type detected, using default")
        return DEFAULT_VIDEO_TYPE

    def _build_prompt(self, content: str, metadata: dict) -> str:
        """构建语义分段 Prompt

        Args:
            content: 字幕内容
            metadata: 视频元数据

        Returns:
            构建好的 Prompt 字符串

        Examples:
            >>> engine = SemanticSegmentationEngine()
            >>> prompt = engine._build_prompt("字幕内容", {"title": "教程"})
            >>> "视频类型" in prompt or "知识类" in prompt
            True
        """
        # 推断视频类型
        video_type = self._infer_video_type(metadata)
        title = metadata.get("title", "") or ""

        # 构建 Prompt
        total_chars = len(content)
        prompt = f"""你是一个专业的视频内容分析师。请将以下字幕内容按语义分段。

【视频信息】
类型：{video_type}
标题：{title}
总字数：{total_chars}字

【字幕内容】
{content}

【分段要求】
1. 建议将内容分为 5-8 个段落（可根据语义完整性灵活调整）
2. 每个段落约 1000-2000 字（为保持语义完整性可以突破此范围）
3. 在话题转换处创建新分段，不要混入不同话题的内容
4. **重要**：保持句子完整性，分段边界应放在自然停顿处（句号。、问号？、感叹号！、省略号……等标点符号后）
   - 基于语义停顿判断，不仅依赖标点符号
   - 保持内容的连贯性和可读性
5. 保持论点完整性，不要在前提-结论结构中间截断
6. 为主题/章节类型的分段生成简洁标题（5-15字），论点/观点类型分段不加标题

【输出格式】
必须返回以下 JSON 格式（不要添加任何额外文本）：
```json
{{
  "video_type": "{video_type}",
  "segments": [
    {{
      "title": "分段标题（如果此分段不需要标题则设为 null）",
      "content": "分段内容",
      "start_index": 0,
      "end_index": 123
    }}
  ],
  "actual_segment_count": 6
}}
```

【注意事项】
- 优先保证语义完整性，不要为了凑数而强行截断
- 标题应简洁准确，如「机器学习概述」「梯度下降原理」
- 如果某个分段无法生成合适的标题（如过渡性内容），将 title 设为 null"""
        return prompt

    def segment(
        self,
        segments: List,
        metadata: dict
    ) -> List[str]:
        """对字幕分段进行语义分段

        Args:
            segments: 字幕分段列表（来自 parser.py 的 Segment 对象）
            metadata: 视频元数据（title, duration 等）

        Returns:
            段落文本列表
        """
        # AI 不可用，直接使用规则分段
        if self.ai_client is None:
            logger.info("AI client unavailable, using fallback segmenter")
            return self.fallback_segmenter.segment(segments)

        # 合并字幕文本
        content = " ".join(seg.text for seg in segments)

        # 构建 Prompt
        prompt = self._build_prompt(
            content=content,
            metadata=metadata
        )

        # 调用 LLM
        try:
            llm_output = self.ai_client.generate(
                prompt,
                max_tokens=8000,
                disable_reasoning=True  # 分段不需要思考模式
            )

            # 解析和验证
            video_segments = self._parse_result(
                llm_output,
                fallback_segments=segments
            )

            # 转换为段落文本
            paragraphs = [seg.content for seg in video_segments]

            logger.info(f"LLM semantic segmentation: {len(paragraphs)} paragraphs")
            return paragraphs

        except Exception as e:
            logger.error(f"LLM segmentation failed: {e}, using fallback")
            return self.fallback_segmenter.segment(segments)

    def _parse_result(
        self,
        llm_output: str,
        fallback_segments: List
    ) -> List[VideoSegment]:
        """解析 LLM 输出并验证

        Args:
            llm_output: LLM 原始输出
            fallback_segments: 降级时使用的原始分段

        Returns:
            VideoSegment 对象列表

        Raises:
            Exception: 解析或验证失败时
        """
        try:
            # 提取 JSON
            data = extract_json_from_llm_response(llm_output)

            # Pydantic 验证
            result = SegmentationResult(**data)

            logger.info(
                f"LLM 分段成功: 类型={result.video_type}, "
                f"分段数={result.actual_segment_count}"
            )

            return result.segments

        except (json.JSONDecodeError, ValidationError) as e:
            logger.error(f"LLM 分段结果验证失败: {e}")
            logger.error(f"原始 LLM 输出: {llm_output[:500]}...")

            # 降级到规则分段
            logger.info("降级到规则分段")
            raise

        except Exception as e:
            logger.error(f"未知错误: {e}")
            raise