"""
索引构建器模块
负责构建向量索引和BM25索引
"""
from pathlib import Path
from typing import List, Optional
from time import time
from loguru import logger

from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
)
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.schema import BaseNode

from src.config import get_config_manager, ConfigManager
from src.storage import BM25StoreManager
from src.storage.strategies import StorageProvider
from src.events import EventBus, IndexBuildEvent
from src.types import IndexResult


class IndexBuilder:
    """
    索引构建器类
    
    使用存储策略模式和事件总线，实现与存储层的解耦
    """
    
    def __init__(self, config: Optional[ConfigManager] = None):
        """
        初始化索引构建器
        
        Args:
            config: 配置实例，如果为None则使用全局配置
        """
        self.config = config or get_config_manager()
        self.storage_provider = StorageProvider.get_instance()
        self.storage_provider.setup(self.config)
        self.bm25_store_manager = BM25StoreManager(config=self.config)
        self.event_bus = EventBus.get_instance()
    
    def load_documents(self, data_dir: str) -> List:
        """
        从目录加载文档
        
        Args:
            data_dir: 文档目录路径
            
        Returns:
            List: 文档列表
        """
        data_path = Path(data_dir)
        
        # 检查目录是否存在
        if not data_path.exists():
            logger.error(f"数据目录不存在: {data_path}")
            return []
        
        # 检查目录是否为空
        files = list(data_path.glob("*"))
        if not files:
            logger.error(f"数据目录为空: {data_path}")
            return []
        
        logger.info("读取文档中...")
        documents = SimpleDirectoryReader(
            input_dir=str(data_path),
            recursive=True,
            required_exts=[".txt", ".md", ".pdf", ".docx"],
        ).load_data()
        
        logger.info(f"成功读取 {len(documents)} 个文档")
        return documents
    
    def parse_documents_to_nodes(self, documents: List) -> List[BaseNode]:
        """
        将文档解析为节点
        
        Args:
            documents: 文档列表
            
        Returns:
            List[BaseNode]: 节点列表
        """
        retrieval_config = self.config.retrieval
        
        logger.info("解析文档为节点...")
        splitter = SentenceSplitter(
            chunk_size=retrieval_config.chunk_size,
            chunk_overlap=retrieval_config.chunk_overlap,
        )
        nodes = splitter.get_nodes_from_documents(documents)
        logger.info(f"生成 {len(nodes)} 个文档节点")
        return nodes
    
    def build_vector_index(self, nodes: List[BaseNode]) -> IndexResult:
        """
        构建向量索引
        
        Args:
            nodes: 文档节点列表
            
        Returns:
            IndexResult: 索引构建结果
        """
        start_time = time()
        
        # 发布构建开始事件
        self.event_bus.publish(
            IndexBuildEvent.build_started(
                nodes_count=len(nodes),
                source="IndexBuilder"
            )
        )
        
        logger.info("构建向量索引中...")
        
        try:
            # 获取当前存储策略
            strategy = self.storage_provider.get_current_strategy()
            
            # 使用策略构建索引
            index = strategy.build_index(nodes, show_progress=True)
            
            if index is None:
                duration = time() - start_time
                # 发布构建失败事件
                self.event_bus.publish(
                    IndexBuildEvent.build_failed(
                        index_type="vector",
                        error="策略构建索引返回None",
                        source="IndexBuilder"
                    )
                )
                return IndexResult.fail(
                    error="构建向量索引失败",
                    index_type="vector",
                    message="策略构建索引返回None"
                )
            
            # 持久化索引
            persist_result = strategy.persist(index)
            if not persist_result.success:
                logger.warning(f"持久化索引失败: {persist_result.error}")
            
            duration = time() - start_time
            
            # 发布构建完成事件
            self.event_bus.publish(
                IndexBuildEvent.build_completed(
                    index_type="vector",
                    nodes_count=len(nodes),
                    duration=duration,
                    source="IndexBuilder"
                )
            )
            
            return IndexResult.ok(
                index=index,
                nodes_count=len(nodes),
                index_type="vector",
                message=f"向量索引已构建 - 后端: {strategy.backend_name}",
                duration=duration
            )
            
        except Exception as e:
            duration = time() - start_time
            logger.error(f"构建向量索引失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            
            # 发布构建失败事件
            self.event_bus.publish(
                IndexBuildEvent.build_failed(
                    index_type="vector",
                    error=str(e),
                    source="IndexBuilder"
                )
            )
            
            return IndexResult.fail(
                error=str(e),
                index_type="vector",
                message="构建向量索引时发生异常"
            )
    
    def build_bm25_index(self, nodes: List[BaseNode]) -> bool:
        """
        构建BM25索引
        
        Args:
            nodes: 文档节点列表
            
        Returns:
            bool: 构建是否成功
        """
        if not self.config.retrieval.use_bm25:
            return True
        
        try:
            logger.info("构建 BM25 索引...")
            self.bm25_store_manager.setup_bm25(nodes=nodes)
            self.bm25_store_manager.persist_bm25()
            logger.info("✅ BM25索引已构建")
            return True
            
        except Exception as e:
            logger.error(f"构建BM25索引失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return False
    
    def build_all_indexes(self, data_dir: str, force_rebuild: bool = False) -> IndexResult:
        """
        构建所有索引（向量索引和BM25索引）
        
        Args:
            data_dir: 文档目录路径
            force_rebuild: 是否强制重建索引
            
        Returns:
            IndexResult: 索引构建结果
        """
        start_time = time()
        
        try:
            logger.info(f"开始构建索引 - 数据目录: {data_dir}")
            
            # 检查是否已存在索引
            strategy = self.storage_provider.get_current_strategy()
            if strategy.exists() and not force_rebuild:
                logger.warning("索引已存在！使用 --force 参数可强制重建")
                return IndexResult.skip(
                    index_type="hybrid",
                    message="索引已存在，使用 --force 参数可强制重建"
                )
            
            # 加载文档
            documents = self.load_documents(data_dir)
            if not documents:
                return IndexResult.fail(
                    error="无法加载文档",
                    index_type="hybrid",
                    message="文档目录为空或不存在"
                )
            
            # 解析为节点
            nodes = self.parse_documents_to_nodes(documents)
            if not nodes:
                return IndexResult.fail(
                    error="无法生成文档节点",
                    index_type="hybrid",
                    message="文档解析失败"
                )
            
            # 构建向量索引
            vector_result = self.build_vector_index(nodes)
            if not vector_result.success:
                return vector_result
            
            # 构建BM25索引
            bm25_success = self.build_bm25_index(nodes)
            
            duration = time() - start_time
            
            # 输出总结
            logger.info("=" * 60)
            logger.info("🎉 所有索引构建完成！")
            logger.info(f"  - 向量索引: ✅ ({len(nodes)} 个节点) - 后端: {strategy.backend_name}")
            if self.config.retrieval.use_bm25:
                logger.info(f"  - BM25索引: {'✅' if bm25_success else '❌'} ({len(nodes)} 个节点)")
            logger.info("=" * 60)
            
            return IndexResult.ok(
                index=vector_result.index,
                nodes_count=len(nodes),
                index_type="hybrid" if self.config.retrieval.use_bm25 else "vector",
                message="所有索引构建完成",
                duration=duration,
                bm25_success=bm25_success
            )
            
        except Exception as e:
            duration = time() - start_time
            logger.error(f"构建索引失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            
            return IndexResult.fail(
                error=str(e),
                index_type="hybrid",
                message="构建索引时发生异常"
            )
