"""
索引加载器模块
负责从存储加载向量索引和BM25索引
"""
from typing import Optional
from time import time
from loguru import logger

from llama_index.core import VectorStoreIndex

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


class IndexLoader:
    """
    索引加载器类
    
    使用存储策略模式和事件总线，实现与存储层的解耦
    """
    
    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_vector_index(self) -> IndexResult:
        """
        加载向量索引
        
        Returns:
            IndexResult: 索引加载结果
        """
        start_time = time()
        
        # 发布加载开始事件
        self.event_bus.publish(
            IndexLoadEvent.load_started(
                index_type="vector",
                source="IndexLoader"
            )
        )
        
        try:
            # 获取当前存储策略
            strategy = self.storage_provider.get_current_strategy()
            
            # 检查索引是否存在
            if not strategy.exists():
                logger.error("向量索引不存在，请先构建索引")
                self.event_bus.publish(
                    IndexLoadEvent.load_failed(
                        index_type="vector",
                        error="索引不存在",
                        source="IndexLoader"
                    )
                )
                return IndexResult.fail(
                    error="向量索引不存在",
                    index_type="vector",
                    message="请先构建索引"
                )
            
            logger.info(f"加载向量索引 - 后端: {strategy.backend_name}")
            
            # 使用策略加载索引
            index = strategy.load_index()
            
            if index is None:
                duration = time() - start_time
                self.event_bus.publish(
                    IndexLoadEvent.load_failed(
                        index_type="vector",
                        error="策略加载索引返回None",
                        source="IndexLoader"
                    )
                )
                return IndexResult.fail(
                    error="加载向量索引失败",
                    index_type="vector",
                    message="策略加载索引返回None"
                )
            
            duration = time() - start_time
            
            # 发布加载完成事件
            self.event_bus.publish(
                IndexLoadEvent.load_completed(
                    index_type="vector",
                    duration=duration,
                    source="IndexLoader"
                )
            )
            
            return IndexResult.ok(
                index=index,
                nodes_count=0,  # 加载时无法直接获取节点数
                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(
                IndexLoadEvent.load_failed(
                    index_type="vector",
                    error=str(e),
                    source="IndexLoader"
                )
            )
            
            return IndexResult.fail(
                error=str(e),
                index_type="vector",
                message="加载向量索引时发生异常"
            )
    
    def load_bm25_index(self, vector_index: Optional[VectorStoreIndex] = None) -> bool:
        """
        加载BM25索引
        
        Args:
            vector_index: 向量索引实例，用于在BM25索引不存在时重建
            
        Returns:
            bool: 加载是否成功
        """
        if not self.config.retrieval.use_bm25:
            return True
        
        try:
            # 尝试从持久化文件加载
            bm25_retriever = self.bm25_store_manager.setup_bm25(nodes=None)
            
            # 如果加载失败且有向量索引，从向量索引的节点重建
            if bm25_retriever is None and vector_index is not None:
                logger.info("BM25 索引加载失败，尝试从向量索引的节点重建...")
                try:
                    # 从向量索引获取所有节点
                    docstore = vector_index.docstore
                    all_nodes = list(docstore.docs.values())
                    
                    if all_nodes:
                        logger.info(f"从向量索引获取到 {len(all_nodes)} 个节点，重建 BM25 索引...")
                        self.bm25_store_manager.setup_bm25(nodes=all_nodes)
                        logger.info("✅ BM25 索引重建成功")
                        return True
                    else:
                        logger.warning("向量索引中没有节点，无法重建 BM25 索引")
                        return False
                except Exception as e:
                    logger.error(f"重建 BM25 索引失败: {e}")
                    return False
            
            return bm25_retriever is not None
            
        except Exception as e:
            logger.error(f"加载BM25索引失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return False
    
    def load_all_indexes(self) -> Optional[VectorStoreIndex]:
        """
        加载所有索引
        
        Returns:
            Optional[VectorStoreIndex]: 向量索引实例，失败返回None
        """
        # 加载向量索引
        vector_index = self.load_vector_index()
        if vector_index is None:
            return None
        
        # 加载BM25索引
        self.load_bm25_index(vector_index)
        
        return vector_index
