"""
索引管理器模块
负责索引的状态查询、删除等管理操作
支持增量索引管理（插入、更新、删除）
"""
from typing import Optional, List, Dict, Any
from loguru import logger
import json
import time
from functools import wraps

from llama_index.core import Document
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, QdrantStorageStrategy


def retry_with_backoff(max_retries: int = 3, initial_delay: float = 1.0, max_delay: float = 30.0):
    """
    带有指数退避的重试装饰器
    
    Args:
        max_retries: 最大重试次数
        initial_delay: 初始延迟时间（秒）
        max_delay: 最大延迟时间（秒）
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            delay = initial_delay
            last_exception = None
            
            for attempt in range(max_retries + 1):
                try:
                    if attempt > 0:
                        logger.info(f"重试 {func.__name__}，第 {attempt} 次尝试，等待 {delay:.1f} 秒...")
                        time.sleep(delay)
                    
                    return func(*args, **kwargs)
                    
                except Exception as e:
                    last_exception = e
                    
                    # 检查是否是超时错误
                    error_msg = str(e).lower()
                    is_timeout = any(keyword in error_msg for keyword in ['timeout', 'timed out', 'read timeout'])
                    
                    if attempt < max_retries and is_timeout:
                        logger.warning(f"{func.__name__} 超时，将在 {delay:.1f} 秒后重试: {e}")
                        # 指数退避
                        delay = min(delay * 2, max_delay)
                    else:
                        # 如果不是超时错误或已达到最大重试次数，则直接抛出异常
                        if attempt == max_retries:
                            logger.error(f"{func.__name__} 失败，已达到最大重试次数 ({max_retries}): {e}")
                        raise
            
            # 理论上不会执行到这里，因为上面已经抛出了异常
            raise last_exception
        
        return wrapper
    return decorator


class IndexManager:
    """索引管理器类"""
    
    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)
    
    def get_index_stats(self, index=None) -> dict:
        """
        获取索引统计信息
        
        Args:
            index: 向量索引实例（可选）
            
        Returns:
            dict: 索引统计信息
        """
        try:
            stats = {}
            
            # 向量索引统计
            if self.config.vector_store.use_qdrant:
                stats.update(self._get_qdrant_stats(index))
            else:
                stats.update(self._get_local_stats(index))
            
            # BM25索引统计
            stats.update(self._get_bm25_stats())
            
            # 检索模式
            stats.update({
                "retrieval_mode": self.config.retrieval.mode,
            })
            
            return stats
            
        except Exception as e:
            logger.error(f"获取索引统计信息失败: {e}")
            return {
                "error": str(e)
            }
    
    def _get_qdrant_stats(self, index) -> dict:
        """获取Qdrant索引统计信息"""
        stats = {
            "storage_type": "qdrant",
        }
        strategy = self.storage_provider.get_current_strategy()
        
        # 确保策略是Qdrant类型
        if not isinstance(strategy, QdrantStorageStrategy):
            stats.update({
                "vector_index_exists": False,
                "message": "当前策略不是Qdrant类型"
            })
            return stats
        
        qdrant_client = strategy.get_qdrant_client()
        collection_name = self.config.vector_store.qdrant_collection_name
        
        # 检查Qdrant客户端是否已初始化
        if not qdrant_client:
            stats.update({
                "vector_index_exists": False,
                "message": "Qdrant 客户端未初始化，请先构建索引"
            })
            return stats
        
        try:
            # 先检查Qdrant服务连接
            try:
                collections = qdrant_client.get_collections()
            except Exception as e:
                logger.error(f"无法连接到 Qdrant 服务: {e}")
                stats.update({
                    "vector_index_exists": False,
                    "connection_error": True,
                    "message": f"无法连接到 Qdrant 服务: {str(e)}"
                })
                return stats
            
            # 检查集合是否存在
            collection_names = [c.name for c in collections.collections]
            if collection_name not in collection_names:
                stats.update({
                    "vector_index_exists": False,
                    "collection_name": collection_name,
                    "message": f"Qdrant 集合 '{collection_name}' 不存在，请先构建索引"
                })
                return stats
            
            # 获取集合信息
            try:
                collection_info = qdrant_client.get_collection(
                    collection_name=collection_name
                )
                # Qdrant API 使用 points_count 而不是 vectors_count
                vectors_count = collection_info.points_count or 0
                
                # 判断索引是否真正可用（有向量数据）
                if vectors_count > 0:
                    stats.update({
                        "vector_index_exists": True,
                        "collection_name": collection_name,
                        "vectors_count": vectors_count,
                        "vector_index_loaded": index is not None,
                        "status": "ready"
                    })
                else:
                    stats.update({
                        "vector_index_exists": False,
                        "collection_name": collection_name,
                        "vectors_count": 0,
                        "message": f"Qdrant 集合 '{collection_name}' 存在但为空，请构建索引"
                    })
            except Exception as e:
                logger.error(f"获取 Qdrant 集合信息失败: {e}")
                stats.update({
                    "vector_index_exists": False,
                    "collection_name": collection_name,
                    "message": f"获取集合信息失败: {str(e)}"
                })
                
        except Exception as e:
            logger.error(f"检查 Qdrant 状态时发生未预期错误: {e}")
            stats.update({
                "vector_index_exists": False,
                "message": f"检查状态失败: {str(e)}"
            })
        
        return stats
    
    def _get_local_stats(self, index) -> dict:
        """获取本地索引统计信息"""
        stats = {}
        vector_storage_dir = self.config.paths.vector_storage_dir
        
        if not vector_storage_dir.exists():
            stats.update({
                "vector_index_exists": False,
                "storage_type": "local",
                "message": "向量索引不存在"
            })
        else:
            index_files = list(vector_storage_dir.glob("*"))
            stats.update({
                "vector_index_exists": True,
                "storage_type": "local",
                "storage_dir": str(vector_storage_dir),
                "file_count": len(index_files),
                "files": [f.name for f in index_files],
                "vector_index_loaded": index is not None,
            })
        return stats
    
    def _get_bm25_stats(self) -> dict:
        """获取BM25索引统计信息"""
        stats = {}
        retrieval_config = self.config.retrieval
        bm25_storage_dir = self.config.paths.bm25_storage_dir
        
        if retrieval_config.use_bm25:
            # 检查必要的BM25索引文件是否存在
            # nodes.json 和 retriever_config.json 是新版本必需的文件
            bm25_exists = (
                bm25_storage_dir.exists()
                and (bm25_storage_dir / "nodes.json").exists()
                and (bm25_storage_dir / "retriever_config.json").exists()
            )
            stats.update({
                "bm25_enabled": True,
                "bm25_index_exists": bm25_exists,
                "bm25_storage_dir": str(bm25_storage_dir),
                "bm25_index_loaded": self.bm25_store_manager.is_loaded(),
                "bm25_language": retrieval_config.bm25_language,
            })
        else:
            stats.update({
                "bm25_enabled": False,
            })
        return stats
    
    def delete_index(self) -> bool:
        """
        删除所有索引
        
        Returns:
            bool: 删除是否成功
        """
        try:
            # 删除向量索引
            if self.config.vector_store.use_qdrant:
                self._delete_qdrant_index()
            else:
                self._delete_local_index()
            
            # 删除BM25索引
            if self.config.retrieval.use_bm25:
                self.bm25_store_manager.delete_bm25_index()
            
            logger.info("=" * 60)
            logger.info("🗑️  所有索引已成功删除")
            logger.info("=" * 60)
            
            return True
            
        except Exception as e:
            logger.error(f"删除索引失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return False
    
    def _delete_qdrant_index(self):
        """删除Qdrant索引"""
        strategy = self.storage_provider.get_current_strategy()
        if isinstance(strategy, QdrantStorageStrategy):
            strategy.delete()
            logger.info("✅ Qdrant 索引已删除")
        else:
            logger.warning("当前策略不是Qdrant类型，无法删除Qdrant索引")
    
    def _delete_local_index(self):
        """删除本地索引"""
        vector_storage_dir = self.config.paths.vector_storage_dir
        
        if vector_storage_dir.exists():
            import shutil
            # 只删除向量索引文件，保留bm25子目录
            for item in vector_storage_dir.iterdir():
                if item.is_file():
                    item.unlink()
                elif item.is_dir() and item.name != "bm25":
                    shutil.rmtree(item)
            logger.info("✅ 本地向量索引已删除")
        else:
            logger.warning("向量索引不存在，无需删除")
    
    # ============================================================
    # 增量索引管理方法（基于LlamaIndex原生支持）
    # ============================================================
    
    @retry_with_backoff(max_retries=3, initial_delay=2.0, max_delay=30.0)
    def insert_documents(
        self,
        index: VectorStoreIndex,
        documents: List[Document]
    ) -> Dict[str, Any]:
        """
        插入新文档到索引（增量添加）
        
        Args:
            index: 向量索引实例
            documents: 文档列表
            
        Returns:
            dict: 操作结果
        """
        try:
            logger.info(f"开始插入 {len(documents)} 个新文档到索引...")
            
            # 使用LlamaIndex原生方法插入节点
            index.insert_nodes(documents)
            
            # 持久化存储
            index.storage_context.persist()
            
            # 同步BM25索引
            if self.config.retrieval.use_bm25:
                self._sync_bm25_index(index)
            
            logger.info(f"✅ 成功插入 {len(documents)} 个文档")
            
            return {
                "success": True,
                "inserted_count": len(documents),
                "message": f"成功插入 {len(documents)} 个文档"
            }
            
        except Exception as e:
            logger.error(f"❌ 插入文档失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return {
                "success": False,
                "error": str(e),
                "inserted_count": 0
            }
    
    @retry_with_backoff(max_retries=3, initial_delay=2.0, max_delay=30.0)
    def update_documents(
        self,
        index: VectorStoreIndex,
        documents: List[Document]
    ) -> Dict[str, Any]:
        """
        更新索引中的文档
        
        Args:
            index: 向量索引实例
            documents: 更新后的文档列表
            
        Returns:
            dict: 操作结果
        """
        try:
            logger.info(f"开始更新 {len(documents)} 个文档...")
            
            # 使用LlamaIndex原生方法刷新文档
            index.refresh_ref_docs(
                documents=documents,
                update_kwargs={"delete_old": True}
            )
            
            # 持久化存储
            index.storage_context.persist()
            
            # 同步BM25索引
            if self.config.retrieval.use_bm25:
                self._sync_bm25_index(index)
            
            logger.info(f"✅ 成功更新 {len(documents)} 个文档")
            
            return {
                "success": True,
                "updated_count": len(documents),
                "message": f"成功更新 {len(documents)} 个文档"
            }
            
        except Exception as e:
            logger.error(f"❌ 更新文档失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return {
                "success": False,
                "error": str(e),
                "updated_count": 0
            }
    
    @retry_with_backoff(max_retries=3, initial_delay=2.0, max_delay=30.0)
    def delete_document(
        self,
        index: VectorStoreIndex,
        doc_id: str
    ) -> Dict[str, Any]:
        """
        从索引中删除文档
        
        Args:
            index: 向量索引实例
            doc_id: 文档ID
            
        Returns:
            dict: 操作结果
        """
        try:
            logger.info(f"开始删除文档: {doc_id}")
            
            # 使用LlamaIndex原生方法删除文档
            index.delete_ref_doc(
                ref_doc_id=doc_id,
                delete_from_docstore=True
            )
            
            # 持久化存储
            index.storage_context.persist()
            
            # 同步BM25索引
            if self.config.retrieval.use_bm25:
                self._sync_bm25_index(index)
            
            logger.info(f"✅ 成功删除文档: {doc_id}")
            
            return {
                "success": True,
                "deleted_doc_id": doc_id,
                "message": f"成功删除文档: {doc_id}"
            }
            
        except Exception as e:
            logger.error(f"❌ 删除文档失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return {
                "success": False,
                "error": str(e),
                "deleted_doc_id": doc_id
            }
    
    def get_document(
        self,
        index: VectorStoreIndex,
        doc_id: str
    ) -> Dict[str, Any]:
        """
        获取索引中的文档
        
        Args:
            index: 向量索引实例
            doc_id: 文档ID
            
        Returns:
            dict: 文档信息
        """
        try:
            doc_store = index.storage_context.docstore
            doc = doc_store.get_document(doc_id)
            
            if doc:
                return {
                    "success": True,
                    "doc_id": doc_id,
                    "text": doc.text,
                    "metadata": doc.metadata,
                    "doc_hash": doc.hash
                }
            else:
                return {
                    "success": False,
                    "error": f"文档 {doc_id} 不存在"
                }
                
        except Exception as e:
            logger.error(f"❌ 获取文档失败: {e}")
            return {
                "success": False,
                "error": str(e)
            }
    
    def list_all_documents(
        self,
        index: VectorStoreIndex
    ) -> Dict[str, Any]:
        """
        列出索引中的所有文档

        Args:
            index: 向量索引实例

        Returns:
            dict: 文档列表信息
        """
        try:
            doc_store = index.storage_context.docstore
            
            # 方法1: 尝试使用docs属性
            all_docs = doc_store.docs
            
            # 如果docs为空，尝试方法2: 从本地文件读取
            if not all_docs:
                try:
                    # 从本地存储文件读取文档
                    docstore_path = self.config.paths.vector_storage_dir / "docstore.json"
                    if docstore_path and docstore_path.exists():
                        with open(docstore_path, 'r', encoding='utf-8') as f:
                            docstore_data = json.load(f)
                        
                        if "docstore/ref_doc_info" in docstore_data:
                            ref_docs = docstore_data["docstore/ref_doc_info"]
                            doc_list = []
                            for doc_id, doc_info in ref_docs.items():
                                # 尝试从docstore/data获取文本
                                text = ""
                                if "docstore/data" in docstore_data:
                                    # 查找对应的节点
                                    node_ids = doc_info.get("node_ids", [])
                                    if node_ids:
                                        node_id = node_ids[0]
                                        if node_id in docstore_data["docstore/data"]:
                                            node_data = docstore_data["docstore/data"][node_id]
                                            if "__data__" in node_data:
                                                text = node_data["__data__"].get("text", "")
                                
                                doc_list.append({
                                    "doc_id": doc_id,
                                    "text_preview": text[:100] + "..." if len(text) > 100 else text,
                                    "metadata": doc_info.get("metadata", {}),
                                    "text_length": len(text)
                                })
                            
                            logger.info(f"✅ 从本地文件读取到 {len(doc_list)} 个文档")
                            
                            return {
                                "success": True,
                                "total_count": len(doc_list),
                                "documents": doc_list
                            }
                except Exception as file_e:
                    logger.warning(f"从本地文件读取文档失败: {file_e}")
            
            # 使用原始方法
            doc_list = []
            for doc_id, doc in all_docs.items():
                doc_list.append({
                    "doc_id": doc_id,
                    "text_preview": doc.text[:100] + "..." if len(doc.text) > 100 else doc.text,
                    "metadata": doc.metadata,
                    "text_length": len(doc.text)
                })
            
            logger.info(f"✅ 索引中共有 {len(doc_list)} 个文档")
            
            return {
                "success": True,
                "total_count": len(doc_list),
                "documents": doc_list
            }
                
        except Exception as e:
            logger.error(f"❌ 列出文档失败: {e}")
            return {
                "success": False,
                "error": str(e),
                "documents": []
            }
    def _sync_bm25_index(self, index: VectorStoreIndex):
        """
        同步BM25索引
        
        Args:
            index: 向量索引实例
        """
        try:
            logger.info("🔄 同步BM25索引...")
            
            # 获取所有文档
            doc_store = index.storage_context.docstore
            all_docs = list(doc_store.docs.values())
            
            if not all_docs:
                logger.warning("⚠️ 没有文档需要同步到BM25索引")
                return
            
            # 重建BM25索引
            self.bm25_store_manager.build_bm25_index(all_docs)
            
            logger.info("✅ BM25索引同步完成")
            
        except Exception as e:
            logger.warning(f"⚠️ BM25索引同步失败: {e}")
