"""
本地存储策略实现
使用文件系统存储向量索引
"""
from typing import Optional, List
from pathlib import Path
from loguru import logger

from llama_index.core import VectorStoreIndex, StorageContext, load_index_from_storage
from llama_index.core.schema import BaseNode

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


class LocalStorageStrategy(StorageStrategy):
    """本地文件系统存储策略"""
    
    def __init__(self, config: StorageConfig):
        """
        初始化本地存储策略
        
        Args:
            config: 存储配置
        """
        if config.backend != StorageBackend.LOCAL:
            raise ValueError("配置的后端类型必须是LOCAL")
        
        super().__init__(config)
        self.persist_dir = Path(config.persist_dir) if config.persist_dir else None
        self.event_bus = EventBus.get_instance()
    
    def initialize(self) -> OperationResult:
        """
        初始化本地存储
        
        Returns:
            操作结果
        """
        try:
            if not self.persist_dir:
                return OperationResult.fail(
                    error="本地存储路径未配置",
                    message="初始化本地存储失败"
                )
            
            # 确保目录存在
            self.persist_dir.mkdir(parents=True, exist_ok=True)
            logger.info(f"本地存储已初始化: {self.persist_dir}")
            
            return OperationResult.ok(
                data=True,
                message=f"本地存储已初始化: {self.persist_dir}"
            )
            
        except Exception as e:
            logger.error(f"初始化本地存储失败: {e}")
            return OperationResult.fail(
                error=str(e),
                message="初始化本地存储失败"
            )
    
    def build_index(self, nodes: List[BaseNode], show_progress: bool = True) -> Optional[VectorStoreIndex]:
        """
        构建向量索引（本地存储）
        
        Args:
            nodes: 文档节点列表
            show_progress: 是否显示进度
            
        Returns:
            向量索引实例或None
        """
        try:
            logger.info("使用本地存储构建向量索引...")
            
            # 直接构建索引（不指定storage_context时使用内存存储）
            index = VectorStoreIndex(
                nodes=nodes,
                show_progress=show_progress
            )
            
            logger.info(f"✅ 向量索引已构建 - 节点数: {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]:
        """
        从本地存储加载向量索引
        
        Returns:
            向量索引实例或None
        """
        try:
            if not self.exists():
                logger.error(f"本地索引不存在: {self.persist_dir}")
                return None
            
            logger.info(f"从本地存储加载向量索引: {self.persist_dir}")
            
            # 创建存储上下文
            self._storage_context = StorageContext.from_defaults(
                persist_dir=str(self.persist_dir)
            )
            
            # 加载索引
            index = load_index_from_storage(self._storage_context)
            
            logger.info("✅ 从本地存储加载向量索引成功")
            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:
        """
        持久化索引到本地存储
        
        Args:
            index: 向量索引实例
            
        Returns:
            操作结果
        """
        try:
            if not self.persist_dir:
                return OperationResult.fail(
                    error="持久化路径未配置",
                    message="持久化索引失败"
                )
            
            logger.info(f"保存向量索引到: {self.persist_dir}")
            
            # 持久化索引
            index.storage_context.persist(persist_dir=str(self.persist_dir))
            
            # 发布存储事件
            self.event_bus.publish(
                StorageEvent.persisted(
                    storage_type="vector",
                    location=str(self.persist_dir),
                    source="LocalStorageStrategy"
                )
            )
            
            logger.info("✅ 向量索引已保存到本地")
            return OperationResult.ok(
                data=True,
                message=f"索引已保存到: {self.persist_dir}"
            )
            
        except Exception as e:
            logger.error(f"持久化索引失败: {e}")
            return OperationResult.fail(
                error=str(e),
                message="持久化索引失败"
            )
    
    def delete(self) -> OperationResult:
        """
        删除本地存储
        
        Returns:
            操作结果
        """
        try:
            if not self.persist_dir or not self.persist_dir.exists():
                return OperationResult.ok(
                    data=True,
                    message="本地存储不存在，无需删除"
                )
            
            logger.info(f"删除本地存储: {self.persist_dir}")
            
            import shutil
            shutil.rmtree(self.persist_dir)
            self.persist_dir.mkdir(parents=True, exist_ok=True)
            
            # 发布存储事件
            self.event_bus.publish(
                StorageEvent.deleted(
                    storage_type="vector",
                    location=str(self.persist_dir),
                    source="LocalStorageStrategy"
                )
            )
            
            logger.info("✅ 本地存储已删除")
            return OperationResult.ok(
                data=True,
                message="本地存储已删除"
            )
            
        except Exception as e:
            logger.error(f"删除本地存储失败: {e}")
            return OperationResult.fail(
                error=str(e),
                message="删除本地存储失败"
            )
    
    def exists(self) -> bool:
        """
        检查本地存储是否存在
        
        Returns:
            是否存在
        """
        if not self.persist_dir:
            return False
        
        # 检查目录是否存在且非空
        if self.persist_dir.exists():
            file_count = len(list(self.persist_dir.glob("*")))
            return file_count > 0
        
        return False
