"""
Qdrant存储策略实现
使用Qdrant向量数据库存储向量索引
支持同步和异步客户端
"""
from typing import Optional, List
from loguru import logger
from qdrant_client import QdrantClient, AsyncQdrantClient

from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.core.schema import BaseNode
from llama_index.vector_stores.qdrant import QdrantVectorStore

from src.types import OperationResult, StorageConfig, StorageBackend
from src.events import EventBus, StorageEvent
from src.utils.patches import apply_all_patches
from .base import StorageStrategy


class QdrantStorageStrategy(StorageStrategy):
    """Qdrant向量数据库存储策略"""
    
    def __init__(self, config: StorageConfig):
        """
        初始化Qdrant存储策略
        
        Args:
            config: 存储配置
        """
        if config.backend != StorageBackend.QDRANT:
            raise ValueError("配置的后端类型必须是QDRANT")
        
        super().__init__(config)
        self.qdrant_client: Optional[QdrantClient] = None
        self.async_qdrant_client: Optional[AsyncQdrantClient] = None
        self.vector_store: Optional[QdrantVectorStore] = None
        self.event_bus = EventBus.get_instance()
        
        # 应用所有兼容性补丁
        apply_all_patches()
    
    def initialize(self) -> OperationResult:
        """
        初始化Qdrant连接（同步和异步客户端）
        
        Returns:
            操作结果
        """
        try:
            logger.info("初始化 Qdrant 向量数据库...")
            
            # 创建同步 Qdrant 客户端
            if self.config.qdrant_url:
                # 使用云服务
                self.qdrant_client = QdrantClient(
                    url=self.config.qdrant_url,
                    api_key=self.config.qdrant_api_key,
                    timeout=60.0,  # 增加超时时间到60秒
                )
                logger.info(f"连接到 Qdrant 云服务: {self.config.qdrant_url[:30]}...")
            else:
                # 使用本地服务
                self.qdrant_client = QdrantClient(
                    host=self.config.qdrant_host,
                    port=self.config.qdrant_port,
                    timeout=60.0,  # 增加超时时间到60秒
                )
                logger.info(
                    f"连接到 Qdrant: "
                    f"{self.config.qdrant_host}:{self.config.qdrant_port}"
                )
            
            # 创建异步 Qdrant 客户端
            if self.config.qdrant_url:
                # 使用云服务
                self.async_qdrant_client = AsyncQdrantClient(
                    url=self.config.qdrant_url,
                    api_key=self.config.qdrant_api_key,
                    timeout=60.0,  # 增加超时时间到60秒
                )
            else:
                # 使用本地服务
                self.async_qdrant_client = AsyncQdrantClient(
                    host=self.config.qdrant_host,
                    port=self.config.qdrant_port,
                    timeout=60.0,  # 增加超时时间到60秒
                )
            
            # 创建向量存储（同时传入同步和异步客户端）
            # 注意：某些版本的 llama-index-vector-stores-qdrant 的 Pydantic 模型
            # 要求 path 字段为必需，即使使用 client 模式也需要传入 path=None
            self.vector_store = QdrantVectorStore(
                client=self.qdrant_client,
                aclient=self.async_qdrant_client,
                collection_name=self.config.qdrant_collection_name,
                path=None,  # 必须显式传入，否则 Pydantic 验证会失败
            )
            
            logger.info(
                f"✅ Qdrant 初始化完成 - 集合: {self.config.qdrant_collection_name}"
            )
            
            return OperationResult.ok(
                data=True,
                message=f"Qdrant已连接 - 集合: {self.config.qdrant_collection_name}"
            )
            
        except Exception as e:
            logger.error(f"初始化 Qdrant 失败: {e}")
            return OperationResult.fail(
                error=str(e),
                message="初始化 Qdrant 失败"
            )
    
    def build_index(self, nodes: List[BaseNode], show_progress: bool = True) -> Optional[VectorStoreIndex]:
        """
        构建向量索引（Qdrant存储）
        
        Args:
            nodes: 文档节点列表
            show_progress: 是否显示进度
        
        Returns:
            向量索引实例或None
        """
        try:
            # 确保Qdrant已初始化
            if not self.vector_store:
                init_result = self.initialize()
                if not init_result.success:
                    return None
            
            logger.info(
                f"使用 Qdrant 构建向量索引 - 集合: {self.config.qdrant_collection_name}"
            )
            
            # 创建存储上下文
            self._storage_context = StorageContext.from_defaults(
                vector_store=self.vector_store
            )
            
            # 构建索引
            index = VectorStoreIndex(
                nodes=nodes,
                storage_context=self._storage_context,
                show_progress=show_progress,
            )
            
            logger.info(f"✅ 向量索引已构建并存储到 Qdrant - 节点数: {len(nodes)}")
            return index
            
        except Exception as e:
            logger.error(f"构建向量索引失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return None
    
    def load_index(self) -> Optional[VectorStoreIndex]:
        """
        从Qdrant加载向量索引
        
        Returns:
            向量索引实例或None
        """
        try:
            # 确保Qdrant已初始化
            if not self.vector_store:
                init_result = self.initialize()
                if not init_result.success:
                    return None
            
            logger.info(
                f"从 Qdrant 加载向量索引 - 集合: {self.config.qdrant_collection_name}"
            )
            
            # 创建存储上下文
            self._storage_context = StorageContext.from_defaults(
                vector_store=self.vector_store
            )
            
            # 从向量存储加载索引
            index = VectorStoreIndex.from_vector_store(
                vector_store=self.vector_store,
                storage_context=self._storage_context,
            )
            
            logger.info("✅ 从 Qdrant 加载向量索引成功")
            return index
            
        except Exception as e:
            logger.error(f"加载向量索引失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return None
    
    def persist(self, index: VectorStoreIndex) -> OperationResult:
        """
        持久化索引到Qdrant（Qdrant自动持久化，此方法仅用于一致性）
        
        Args:
            index: 向量索引实例
        
        Returns:
            操作结果
        """
        try:
            # Qdrant 在写入时自动持久化，无需额外操作
            logger.info("Qdrant 自动持久化索引")
            
            # 发布存储事件
            self.event_bus.publish(
                StorageEvent.persisted(
                    storage_type="vector",
                    location=f"qdrant://{self.config.qdrant_collection_name}",
                    source="QdrantStorageStrategy"
                )
            )
            
            return OperationResult.ok(
                data=True,
                message=f"索引已存储到 Qdrant 集合: {self.config.qdrant_collection_name}"
            )
            
        except Exception as e:
            logger.error(f"持久化索引失败: {e}")
            return OperationResult.fail(
                error=str(e),
                message="持久化索引失败"
            )
    
    def delete(self) -> OperationResult:
        """
        删除Qdrant集合
        
        Returns:
            操作结果
        """
        try:
            if not self.qdrant_client:
                init_result = self.initialize()
                if not init_result.success:
                    return OperationResult.fail(
                        error="Qdrant客户端未初始化",
                        message="删除Qdrant集合失败"
                    )
            
            collection_name = self.config.qdrant_collection_name
            logger.info(f"删除 Qdrant 集合: {collection_name}")
            
            self.qdrant_client.delete_collection(collection_name=collection_name)
            
            # 发布存储事件
            self.event_bus.publish(
                StorageEvent.deleted(
                    storage_type="vector",
                    location=f"qdrant://{collection_name}",
                    source="QdrantStorageStrategy"
                )
            )
            
            logger.info(f"✅ Qdrant 集合已删除: {collection_name}")
            return OperationResult.ok(
                data=True,
                message=f"Qdrant集合已删除: {collection_name}"
            )
            
        except Exception as e:
            logger.warning(f"删除 Qdrant 集合时出错（可能不存在）: {e}")
            return OperationResult.fail(
                error=str(e),
                message="删除Qdrant集合失败"
            )
    
    def exists(self) -> bool:
        """
        检查Qdrant集合是否存在
        
        Returns:
            是否存在
        """
        try:
            if not self.qdrant_client:
                init_result = self.initialize()
                if not init_result.success:
                    return False
            
            collections = self.qdrant_client.get_collections().collections
            collection_names = [c.name for c in collections]
            return self.config.qdrant_collection_name in collection_names
            
        except Exception as e:
            logger.warning(f"检查 Qdrant 集合是否存在时出错: {e}")
            return False
    
    def get_vector_store(self) -> Optional[QdrantVectorStore]:
        """
        获取Qdrant向量存储实例
        
        Returns:
            向量存储实例或None
        """
        return self.vector_store
    
    def get_qdrant_client(self) -> Optional[QdrantClient]:
        """
        获取Qdrant客户端
        
        Returns:
            Qdrant客户端或None
        """
        return self.qdrant_client
