"""
Q&A 自动拆分分块策略

使用 LLM 将文档内容自动拆分为问答对。
每个分段的 content 存储 question，metadata 中存储 answer。
检索时用 question 做向量匹配，返回 answer 作为上下文。
"""
import json
import logging
from typing import Dict, Any, List, Optional

from .base import BaseChunker, ChunkResult

logger = logging.getLogger(__name__)

# Q&A 拆分的系统提示词
QA_SYSTEM_PROMPT = """你是一个专业的知识库问答对生成助手。请根据给定的文本内容，生成高质量的问答对（Q&A pairs）。

要求：
1. 问题应该是用户可能会问的自然语言问题
2. 答案应该准确、完整，直接来源于原文
3. 每个问答对应该覆盖文本中的一个独立知识点
4. 问题要具体明确，避免过于宽泛
5. 答案要简洁但完整，包含必要的上下文

请以 JSON 数组格式输出，每个元素包含 question 和 answer 字段：
```json
[
  {"question": "问题1", "answer": "答案1"},
  {"question": "问题2", "answer": "答案2"}
]
```

只输出 JSON 数组，不要输出其他内容。"""


class QAChunker(BaseChunker):
    """
    Q&A 自动拆分分块器

    使用 LLM 将文本拆分为问答对。
    需要在初始化时传入 LLM 调用函数。
    """

    def __init__(
        self,
        chunk_size: int = 500,
        chunk_overlap: int = 50,
        separator: Optional[str] = None,
        llm_caller: Optional[Any] = None,
    ):
        super().__init__(chunk_size, chunk_overlap, separator)
        self._llm_caller = llm_caller

    def chunk(self, text: str, metadata: Dict[str, Any] = None) -> List[ChunkResult]:
        """
        同步分块（Q&A 模式不支持同步调用，返回空列表）
        请使用 chunk_async 方法。
        """
        logger.warning("QAChunker.chunk() 不支持同步调用，请使用 chunk_async()")
        return []

    async def chunk_async(self, text: str, metadata: Dict[str, Any] = None) -> List[ChunkResult]:
        """
        异步分块：使用 LLM 将文本拆分为 Q&A 对

        Args:
            text: 原始文本
            metadata: 文档元数据

        Returns:
            分块结果列表，每个 ChunkResult 的 content 为 question，
            metadata 中包含 answer 和 chunk_mode='qa'
        """
        if not text or not text.strip():
            return []

        metadata = metadata or {}
        text = self._clean_text(text)

        # 如果文本太长，先按段落粗分再逐段生成 Q&A
        max_input_size = self.chunk_size * 8  # LLM 输入上限
        if len(text) > max_input_size:
            segments = self._split_for_qa(text, max_input_size)
        else:
            segments = [text]

        all_results = []
        position = 0

        for segment in segments:
            qa_pairs = await self._generate_qa_pairs(segment)
            for qa in qa_pairs:
                question = qa.get('question', '').strip()
                answer = qa.get('answer', '').strip()
                if not question or not answer:
                    continue

                all_results.append(ChunkResult(
                    content=question,
                    position=position,
                    metadata={
                        **metadata,
                        'answer': answer,
                        'chunk_mode': 'qa',
                    },
                ))
                position += 1

        logger.info(f"Q&A 拆分完成: {len(all_results)} 个问答对")
        return all_results

    async def _generate_qa_pairs(self, text: str) -> List[Dict[str, str]]:
        """调用 LLM 生成 Q&A 对"""
        if not self._llm_caller:
            logger.error("QAChunker: 未配置 LLM 调用函数")
            return []

        try:
            user_prompt = f"请根据以下文本生成问答对：\n\n{text}"

            response_text = await self._llm_caller(
                system_prompt=QA_SYSTEM_PROMPT,
                user_prompt=user_prompt,
            )

            if not response_text:
                return []

            # 解析 JSON 响应
            return self._parse_qa_response(response_text)

        except Exception as e:
            logger.error(f"Q&A 生成失败: {e}")
            return []

    @staticmethod
    def _parse_qa_response(response_text: str) -> List[Dict[str, str]]:
        """解析 LLM 返回的 Q&A JSON"""
        try:
            # 尝试直接解析
            result = json.loads(response_text)
            if isinstance(result, list):
                return [
                    item for item in result
                    if isinstance(item, dict) and 'question' in item and 'answer' in item
                ]
        except json.JSONDecodeError:
            pass

        # 尝试提取 JSON 代码块
        import re
        json_match = re.search(r'```(?:json)?\s*\n?(.*?)\n?```', response_text, re.DOTALL)
        if json_match:
            try:
                result = json.loads(json_match.group(1))
                if isinstance(result, list):
                    return [
                        item for item in result
                        if isinstance(item, dict) and 'question' in item and 'answer' in item
                    ]
            except json.JSONDecodeError:
                pass

        # 尝试找到 [ ... ] 部分
        bracket_match = re.search(r'\[.*\]', response_text, re.DOTALL)
        if bracket_match:
            try:
                result = json.loads(bracket_match.group(0))
                if isinstance(result, list):
                    return [
                        item for item in result
                        if isinstance(item, dict) and 'question' in item and 'answer' in item
                    ]
            except json.JSONDecodeError:
                pass

        logger.warning(f"无法解析 Q&A 响应: {response_text[:200]}")
        return []

    def _split_for_qa(self, text: str, max_size: int) -> List[str]:
        """将长文本按段落分割为适合 LLM 处理的片段"""
        paragraphs = text.split('\n\n')
        segments = []
        current = ""

        for para in paragraphs:
            if current and len(current) + len(para) + 2 > max_size:
                segments.append(current.strip())
                current = para
            else:
                current = current + "\n\n" + para if current else para

        if current.strip():
            segments.append(current.strip())

        return segments