"""
文档索引服务

负责文档处理管道：文本提取 → 分块 → 向量化 → 入库
分段数据存入业务数据库，向量数据存入 Qdrant 向量数据库
"""
import hashlib
import logging
import math
from datetime import datetime
from typing import List, Optional, Tuple

from sqlalchemy import select, func, delete
from sqlalchemy.ext.asyncio import AsyncSession

from ai_platform.knowledge.models import KnowledgeBase, KnowledgeDocument, KnowledgeSegment
from ai_platform.knowledge.chunking import get_chunker
from ai_platform.knowledge.services.embedding_service import EmbeddingService
from ai_platform.knowledge.vector_store import get_vector_store, VectorPoint

logger = logging.getLogger(__name__)

# 向量化批次大小
EMBEDDING_BATCH_SIZE = 50


class IndexingService:
    """
    文档索引服务
    
    处理管道：
    1. 从文件管理系统提取文本内容
    2. 按知识库配置的策略分块
    3. 调用 Embedding 模型向量化
    4. 将分段写入业务数据库，向量写入 Qdrant
    """

    def __init__(self, db: AsyncSession):
        self._db = db
        self._embedding_service = EmbeddingService(db)
        self._vector_store = get_vector_store()

    async def index_document(
        self,
        knowledge_base_id: str,
        document_id: str,
    ) -> Tuple[int, int]:
        """
        索引单个文档
        
        Args:
            knowledge_base_id: 知识库 ID
            document_id: 文档 ID
            
        Returns:
            (segment_count, token_count) 分段数和 Token 数
        """
        # 1. 获取知识库配置
        kb_result = await self._db.execute(
            select(KnowledgeBase).where(
                KnowledgeBase.id == knowledge_base_id,
                KnowledgeBase.is_deleted == False
            )
        )
        kb = kb_result.scalar_one_or_none()
        if not kb:
            raise ValueError(f'知识库不存在: {knowledge_base_id}')

        if not kb.embedding_model_id:
            raise ValueError('知识库未配置 Embedding 模型')

        # 2. 获取文档
        doc_result = await self._db.execute(
            select(KnowledgeDocument).where(
                KnowledgeDocument.id == document_id,
                KnowledgeDocument.is_deleted == False
            )
        )
        doc = doc_result.scalar_one_or_none()
        if not doc:
            raise ValueError(f'文档不存在: {document_id}')

        # 更新状态为 indexing
        doc.status = 'indexing'
        doc.indexing_started_at = datetime.now()
        doc.error_message = None
        await self._db.commit()

        try:
            from ai_platform.knowledge.services.indexing_progress_service import IndexingProgressService

            # 3. 提取文本
            await IndexingProgressService.publish(
                knowledge_base_id, document_id, step='extracting', progress=0.1,
                message='正在提取文本内容...', document_name=doc.name,
            )
            text_content = await self._extract_text(doc.file_id)
            if not text_content or not text_content.strip():
                raise ValueError('文档内容为空，无法索引')

            # 3.5 预处理/清洗
            await IndexingProgressService.publish(
                knowledge_base_id, document_id, step='cleaning', progress=0.2,
                message='正在预处理/清洗文本...', document_name=doc.name,
            )
            from ai_platform.knowledge.services.cleaning_service import CleaningService
            text_content = CleaningService.clean(text_content, kb.process_rules)

            # 计算内容哈希（用于去重）
            content_hash = hashlib.md5(text_content.encode('utf-8')).hexdigest()
            doc.content_hash = content_hash

            # 4. 分块
            await IndexingProgressService.publish(
                knowledge_base_id, document_id, step='chunking', progress=0.3,
                message='正在分块...', document_name=doc.name,
            )
            chunk_strategy = kb.chunk_strategy or 'recursive'
            chunk_kwargs = {}

            # Q&A 模式需要 LLM 调用函数
            if chunk_strategy == 'qa':
                chunk_kwargs['llm_caller'] = self._create_llm_caller(kb)

            chunker = get_chunker(
                strategy=chunk_strategy,
                chunk_size=kb.chunk_size or 500,
                chunk_overlap=kb.chunk_overlap or 50,
                separator=kb.separator,
                **chunk_kwargs,
            )

            doc_metadata = {
                'document_id': document_id,
                'document_name': doc.name,
                'file_type': doc.file_type,
            }

            # Q&A 模式使用异步分块
            if chunk_strategy == 'qa' and hasattr(chunker, 'chunk_async'):
                chunks = await chunker.chunk_async(text_content, metadata=doc_metadata)
            else:
                chunks = chunker.chunk(text_content, metadata=doc_metadata)

            if not chunks:
                raise ValueError('文档分块结果为空')

            # 5. 删除旧的分段（重新索引场景）
            await self._delete_document_segments(document_id, knowledge_base_id)

            # 判断索引模式
            is_economy = (kb.indexing_technique == 'economy')

            # 6. 确保 Qdrant collection 存在（经济模式跳过）
            if not is_economy:
                # 自动检测并修正 embedding 维度
                try:
                    real_dim = await self._embedding_service.get_embedding_dimensions(kb.embedding_model_id)
                    if real_dim != kb.embedding_dimensions:
                        logger.info(f'修正 embedding 维度: {kb.embedding_dimensions} → {real_dim}')
                        kb.embedding_dimensions = real_dim
                        await self._db.commit()
                except Exception as e:
                    logger.warning(f'自动检测 embedding 维度失败: {e}')

                vector_size = kb.embedding_dimensions or 1536
                await self._vector_store.ensure_collection(knowledge_base_id, vector_size)

            # 7. 入库（分批处理）
            segment_count = 0
            total_token_count = 0
            total_char_count = 0
            failed_embedding_count = 0

            total_batches = math.ceil(len(chunks) / EMBEDDING_BATCH_SIZE)

            for batch_idx in range(total_batches):
                start = batch_idx * EMBEDDING_BATCH_SIZE
                end = min(start + EMBEDDING_BATCH_SIZE, len(chunks))
                batch_chunks = chunks[start:end]

                # 高质量模式：批量向量化；经济模式：跳过
                if is_economy:
                    embeddings = [None] * len(batch_chunks)
                else:
                    batch_texts = [c.content for c in batch_chunks]
                    try:
                        embeddings = await self._embedding_service.embed_texts(
                            model_id=kb.embedding_model_id,
                            texts=batch_texts,
                            dimensions=kb.embedding_dimensions if kb.embedding_dimensions else None,
                        )
                    except Exception as e:
                        logger.error(f'向量化批次 {batch_idx + 1}/{total_batches} 失败: {e}')
                        embeddings = [None] * len(batch_texts)
                        failed_embedding_count += len(batch_texts)

                # 创建分段记录（业务数据库）+ 收集向量点（Qdrant）
                vector_points = []
                for i, chunk in enumerate(batch_chunks):
                    embedding = embeddings[i] if i < len(embeddings) else None
                    char_count = len(chunk.content)
                    token_count = self._estimate_tokens(chunk.content)
                    word_count = self._count_words(chunk.content)

                    # 自动提取关键词（高质量和经济模式均提取，增强全文检索）
                    keywords = chunk.metadata.get('keywords')
                    if not keywords:
                        keywords = self._extract_keywords(chunk.content)

                    # 经济模式下 embedding_status 标记为 'skipped'
                    if is_economy:
                        emb_status = 'skipped'
                    else:
                        emb_status = 'completed' if embedding else 'failed'

                    segment = KnowledgeSegment(
                        knowledge_base_id=knowledge_base_id,
                        document_id=document_id,
                        position=start + i,
                        content=chunk.content,
                        answer=chunk.metadata.get('answer'),
                        token_count=token_count,
                        char_count=char_count,
                        word_count=word_count,
                        page_number=chunk.metadata.get('page_number'),
                        keywords=keywords,
                        extra_metadata=chunk.metadata,
                        embedding_status=emb_status,
                        enabled=True,
                    )
                    self._db.add(segment)
                    await self._db.flush()

                    # 收集向量点，稍后批量写入 Qdrant（经济模式跳过）
                    if embedding and not is_economy:
                        vector_points.append(VectorPoint(
                            id=str(segment.id),
                            vector=embedding,
                            payload={
                                'document_id': document_id,
                                'knowledge_base_id': knowledge_base_id,
                                'position': start + i,
                            },
                        ))

                    segment_count += 1
                    total_token_count += token_count
                    total_char_count += char_count

                # 提交业务数据库
                await self._db.commit()

                # 批量写入 Qdrant（经济模式跳过）
                if vector_points:
                    await self._vector_store.upsert(knowledge_base_id, vector_points)

                step_label = '关键词提取中' if is_economy else '向量化中'
                batch_progress = 0.3 + 0.6 * (end / len(chunks))
                await IndexingProgressService.publish(
                    knowledge_base_id, document_id, step='vectorizing',
                    progress=batch_progress,
                    message=f'{step_label} {end}/{len(chunks)}',
                    document_name=doc.name,
                )
                logger.info(f'文档 {doc.name} 索引进度: {end}/{len(chunks)}')

            # 7. 更新文档状态
            if not is_economy and failed_embedding_count > 0:
                if failed_embedding_count >= segment_count:
                    doc.status = 'failed'
                    doc.error_message = f'所有 {segment_count} 个分段向量化失败'
                else:
                    doc.status = 'completed'
                    doc.error_message = f'{failed_embedding_count}/{segment_count} 个分段向量化失败'
            else:
                doc.status = 'completed'
            doc.segment_count = segment_count
            doc.token_count = total_token_count
            doc.char_count = total_char_count
            doc.indexing_completed_at = datetime.now()

            # 8. 更新知识库统计
            await self._update_kb_stats(knowledge_base_id)

            await self._db.commit()

            await IndexingProgressService.publish(
                knowledge_base_id, document_id, step='completed', progress=1.0,
                message=f'索引完成: {segment_count} 个分段',
                document_name=doc.name,
            )
            logger.info(f'文档 {doc.name} 索引完成: {segment_count} 个分段, {total_token_count} tokens')
            return segment_count, total_token_count

        except Exception as e:
            logger.exception(f'文档索引失败: {e}')
            doc.status = 'failed'
            doc.error_message = str(e)[:500]
            await self._db.commit()
            await IndexingProgressService.publish(
                knowledge_base_id, document_id, step='failed', progress=0.0,
                message='索引失败', document_name=doc.name,
                error=str(e)[:200],
            )
            raise

    async def reindex_document(self, knowledge_base_id: str, document_id: str) -> Tuple[int, int]:
        """重新索引文档（删除旧分段后重新处理）"""
        return await self.index_document(knowledge_base_id, document_id)

    async def index_segment(
        self,
        knowledge_base_id: str,
        segment_id: str,
    ) -> bool:
        """
        为单个分段生成向量（用于手动添加或更新分段后）
        """
        kb_result = await self._db.execute(
            select(KnowledgeBase).where(
                KnowledgeBase.id == knowledge_base_id,
                KnowledgeBase.is_deleted == False
            )
        )
        kb = kb_result.scalar_one_or_none()
        if not kb or not kb.embedding_model_id:
            return False

        seg_result = await self._db.execute(
            select(KnowledgeSegment).where(
                KnowledgeSegment.id == segment_id,
                KnowledgeSegment.is_deleted == False
            )
        )
        segment = seg_result.scalar_one_or_none()
        if not segment:
            return False

        try:
            embedding = await self._embedding_service.embed_text(
                model_id=kb.embedding_model_id,
                text=segment.content,
                dimensions=kb.embedding_dimensions if kb.embedding_dimensions else None,
            )

            # 确保 collection 存在
            vector_size = kb.embedding_dimensions or 1536
            await self._vector_store.ensure_collection(knowledge_base_id, vector_size)

            # 写入 Qdrant
            await self._vector_store.upsert(knowledge_base_id, [VectorPoint(
                id=str(segment_id),
                vector=embedding,
                payload={
                    'document_id': str(segment.document_id),
                    'knowledge_base_id': knowledge_base_id,
                    'position': segment.position or 0,
                },
            )])

            segment.embedding_status = 'completed'
            await self._db.commit()
            return True

        except Exception as e:
            logger.error(f'分段向量化失败: {e}')
            segment.embedding_status = 'failed'
            await self._db.commit()
            return False

    async def _extract_text(self, file_id: str) -> str:
        """从文件管理系统提取文本内容（启用 OCR 支持图片和扫描版 PDF）"""
        from core.file_manager.service import FileManagerService

        text_content = await FileManagerService.get_file_text_content(
            self._db, file_id, enable_ocr=True
        )
        if not text_content:
            raise ValueError('无法提取文件文本内容')
        return text_content

    async def _delete_document_segments(self, document_id: str, knowledge_base_id: str):
        """删除文档的所有分段（业务数据库 + Qdrant）"""
        # 先从 Qdrant 删除该文档的所有向量
        try:
            await self._vector_store.delete_by_filter(
                knowledge_base_id,
                filter_conditions={'document_id': document_id},
            )
        except Exception as e:
            logger.warning(f'从 Qdrant 删除文档向量失败: {e}')

        # 再从业务数据库删除分段记录
        await self._db.execute(
            delete(KnowledgeSegment).where(
                KnowledgeSegment.document_id == document_id
            )
        )

    async def _update_kb_stats(self, knowledge_base_id: str):
        """更新知识库统计信息"""
        # 文档数
        doc_count_result = await self._db.execute(
            select(func.count()).select_from(KnowledgeDocument).where(
                KnowledgeDocument.knowledge_base_id == knowledge_base_id,
                KnowledgeDocument.is_deleted == False,
            )
        )
        doc_count = doc_count_result.scalar() or 0

        # 分段数和 Token 数
        seg_stats = await self._db.execute(
            select(
                func.count(),
                func.coalesce(func.sum(KnowledgeSegment.token_count), 0),
                func.coalesce(func.sum(KnowledgeSegment.char_count), 0),
            ).where(
                KnowledgeSegment.knowledge_base_id == knowledge_base_id,
                KnowledgeSegment.is_deleted == False,
            )
        )
        row = seg_stats.one()
        seg_count = row[0] or 0
        total_tokens = row[1] or 0
        total_chars = row[2] or 0

        # 更新知识库
        kb_result = await self._db.execute(
            select(KnowledgeBase).where(KnowledgeBase.id == knowledge_base_id)
        )
        kb = kb_result.scalar_one_or_none()
        if kb:
            kb.document_count = doc_count
            kb.segment_count = seg_count
            kb.total_token_count = total_tokens
            kb.total_char_count = total_chars

    def _create_llm_caller(self, kb: KnowledgeBase):
        """
        创建 Q&A 分块所需的 LLM 调用函数

        使用知识库所属应用中配置的第一个 chat 类型模型。
        """
        db = self._db

        async def llm_caller(system_prompt: str, user_prompt: str) -> str:
            from ai_platform.services.llm_service import LLMService

            # 查找可用的 chat 模型
            from ai_platform.models import LLMModel
            model_result = await db.execute(
                select(LLMModel).where(
                    LLMModel.model_type == 'chat',
                    LLMModel.is_active == True,
                    LLMModel.is_deleted == False,
                ).limit(1)
            )
            model = model_result.scalar_one_or_none()
            if not model:
                raise ValueError('未找到可用的 chat 模型，无法进行 Q&A 拆分')

            llm_service = LLMService(db)
            messages = []
            if system_prompt:
                messages.append({'role': 'system', 'content': system_prompt})
            messages.append({'role': 'user', 'content': user_prompt})

            response = await llm_service.chat_async(
                model_id=str(model.id),
                messages=messages,
                temperature=0.3,
                max_tokens=4096,
            )
            return response.content or ''

        return llm_caller

    @staticmethod
    def _extract_keywords(text: str, max_keywords: int = 10) -> List[str]:
        """
        从文本中提取关键词（经济模式使用）
        
        使用简单的 TF 统计提取高频词，无需外部依赖。
        """
        import re
        if not text:
            return []

        # 中文分词（简单按标点和空格分割）
        # 提取中文词组（2-4字）和英文单词
        chinese_words = re.findall(r'[\u4e00-\u9fff]{2,4}', text)
        english_words = [w.lower() for w in re.findall(r'[a-zA-Z]{3,}', text)]

        all_words = chinese_words + english_words

        # 停用词（简单列表）
        stop_words = {
            '的', '了', '在', '是', '我', '有', '和', '就', '不', '人', '都', '一',
            '一个', '上', '也', '很', '到', '说', '要', '去', '你', '会', '着',
            '没有', '看', '好', '自己', '这', '他', '她', '它', '我们', '他们',
            '可以', '这个', '那个', '什么', '如果', '因为', '所以', '但是', '而且',
            'the', 'and', 'for', 'are', 'but', 'not', 'you', 'all', 'can',
            'had', 'her', 'was', 'one', 'our', 'out', 'has', 'have', 'been',
            'this', 'that', 'with', 'from', 'they', 'will', 'would', 'there',
        }

        # 词频统计
        word_freq = {}
        for word in all_words:
            if word in stop_words or len(word) < 2:
                continue
            word_freq[word] = word_freq.get(word, 0) + 1

        # 按频率排序取 top
        sorted_words = sorted(word_freq.items(), key=lambda x: x[1], reverse=True)
        return [w for w, _ in sorted_words[:max_keywords]]

    @staticmethod
    def _count_words(text: str) -> int:
        """计算词数（中文按字计算，英文按空格分词）"""
        if not text:
            return 0
        import re
        # 中文字符数
        chinese_chars = len(re.findall(r'[\u4e00-\u9fff]', text))
        # 英文单词数
        english_words = len(re.findall(r'[a-zA-Z]+', text))
        return chinese_chars + english_words

    @staticmethod
    def _estimate_tokens(text: str) -> int:
        """粗略估算文本的 Token 数"""
        if not text:
            return 0
        # 中文约 1 字 = 1.5 token，英文约 4 字符 = 1 token
        # 简单混合估算
        chinese_chars = sum(1 for c in text if '\u4e00' <= c <= '\u9fff')
        other_chars = len(text) - chinese_chars
        return int(chinese_chars * 1.5 + other_chars / 4)
