"""
# flake8: noqa
sharded_faiss_index.py
----------------------
Sharded FAISS index implementation for distributed vector storage and search.

Provides automatic shard routing, parallel cross-shard search, and result merging
for scaling to millions of vectors.

Author: HIEROMEM Team
"""

from __future__ import annotations
from typing import List, Tuple, Optional, Dict, Any
import json
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
import numpy as np
from dataclasses import dataclass

from .faiss_index import FaissIndex
from .index_shard_manager import IndexShardManager


@dataclass
class SearchResult:
    """Single search result with episode ID, similarity, and metadata."""
    episode_id: str
    similarity: float
    metadata: Dict[str, Any]


class ShardedFaissIndex:
    """
    Distributed FAISS index using consistent hashing for shard selection.

    Features:
    - Automatic shard routing based on episode ID
    - Parallel cross-shard search
    - Result merging and re-ranking
    - Lazy shard loading for memory efficiency
    """

    def __init__(
        self,
        shard_manager: IndexShardManager,
        dimension: int,
        index_type: str = "Flat",
        metric: str = "ip",
        max_workers: int = 4,
        debug: bool = False
    ):
        """
        Initialize sharded FAISS index.

        Args:
            shard_manager: Shard manager instance
            dimension: Vector dimension
            index_type: FAISS index type ("Flat", "IVF", etc.)
            metric: Distance metric ("cosine", "l2", "ip")
            max_workers: Max parallel workers for cross-shard search
            debug: Enable debug logging
        """
        self.shard_manager = shard_manager
        self.dimension = dimension
        self.index_type = index_type
        self.metric = metric
        self.max_workers = max_workers
        self.debug = debug

        # Cache for loaded shard indices
        self._shard_cache: Dict[int, FaissIndex] = {}
        # Cache for loaded shard metadata
        self._metadata_cache: Dict[int, Dict[str, Dict[str, Any]]] = {}

        self._log(
            f"Initialized sharded index with {shard_manager.get_num_shards()} shards")

    def _log(self, message: str) -> None:
        """Log debug messages if debug mode is enabled."""
        if self.debug:
            print(f"[ShardedFaissIndex] {message}")

    def _get_shard_index(self, shard_id: int) -> FaissIndex:
        """
        Get or load a shard index.

        Args:
            shard_id: Shard ID

        Returns:
            FaissIndex instance
        """
        if shard_id in self._shard_cache:
            return self._shard_cache[shard_id]

        # Load or create shard index
        shard_path = self.shard_manager.get_shard_path(shard_id)

        try:
            if not shard_path.exists():
                raise FileNotFoundError
            # Try to load existing index
            index = FaissIndex.load(shard_path)
            self._log(f"Loaded shard {shard_id} from {shard_path}")
        except (FileNotFoundError, ValueError, RuntimeError):
            # Create new index with correct parameters
            # Use Flat index type to avoid training requirement
            index = FaissIndex(
                dim=self.dimension,
                index_type="hnsw",  # HNSW doesn't require training
                metric=self.metric
            )
            self._log(f"Created new shard {shard_id}")

        self._shard_cache[shard_id] = index
        return index

    def _get_shard_metadata(self, shard_id: int) -> Dict[str, Dict[str, Any]]:
        """Get or load shard metadata."""
        if shard_id in self._metadata_cache:
            return self._metadata_cache[shard_id]

        shard_path = self.shard_manager.get_shard_path(shard_id)
        meta_path = shard_path.with_suffix(".meta.json")
        
        metadata = {}
        if meta_path.exists():
            try:
                with open(meta_path, "r") as f:
                    metadata = json.load(f)
            except Exception as e:
                self._log(f"Failed to load metadata for shard {shard_id}: {e}")
        
        self._metadata_cache[shard_id] = metadata
        return metadata

    def add(
        self,
        episode_id: str,
        vector: np.ndarray,
        metadata: Optional[Dict[str, Any]] = None
    ) -> int:
        """
        Add a vector to the appropriate shard.

        Args:
            episode_id: Episode identifier
            vector: Vector to add
            metadata: Optional metadata

        Returns:
            Shard ID where vector was added
        """
        # Determine shard using consistent hashing
        shard_id = self.shard_manager.get_shard_for_key(episode_id)

        # Get shard index
        shard_index = self._get_shard_index(shard_id)

        # Add to shard (FaissIndex.add only takes episode_id and vector)
        shard_index.add(episode_id, vector)

        # Store metadata separately
        if metadata:
            shard_meta = self._get_shard_metadata(shard_id)
            shard_meta[episode_id] = metadata

        # Update shard stats
        self._update_shard_stats(shard_id, shard_index)

        self._log(f"Added {episode_id} to shard {shard_id}")
        return shard_id

    def add_batch(
        self,
        episode_ids: List[str],
        vectors: np.ndarray,
        metadata_list: Optional[List[Dict[str, Any]]] = None
    ) -> Dict[int, int]:
        """
        Add multiple vectors in batch, routing to appropriate shards.

        Args:
            episode_ids: List of episode IDs
            vectors: Array of vectors (n_vectors, dimension)
            metadata_list: Optional list of metadata dicts

        Returns:
            Dictionary mapping shard_id to count of vectors added
        """
        if metadata_list is None:
            metadata_list = [None] * len(episode_ids)

        # Group by shard
        shard_batches: Dict[int,
                            List[Tuple[str, np.ndarray, Optional[Dict]]]] = {}

        for episode_id, vector, metadata in zip(
                episode_ids, vectors, metadata_list):
            shard_id = self.shard_manager.get_shard_for_key(episode_id)
            if shard_id not in shard_batches:
                shard_batches[shard_id] = []
            shard_batches[shard_id].append((episode_id, vector, metadata))

        # Add to each shard
        shard_counts = {}
        for shard_id, batch in shard_batches.items():
            shard_index = self._get_shard_index(shard_id)

            batch_ids = [item[0] for item in batch]
            batch_vectors = np.array([item[1] for item in batch])
            batch_metadata = [item[2] for item in batch]

            for episode_id, vector in zip(batch_ids, batch_vectors):
                shard_index.add(episode_id, vector)
            
            # Update metadata
            shard_meta = self._get_shard_metadata(shard_id)
            for episode_id, meta in zip(batch_ids, batch_metadata):
                if meta:
                    shard_meta[episode_id] = meta

            shard_counts[shard_id] = len(batch)
            self._update_shard_stats(shard_id, shard_index)

        self._log(
            f"Added {len(episode_ids)} vectors across {len(shard_batches)} shards")
        return shard_counts

    def search(
        self,
        query_vector: np.ndarray,
        k: int = 5,
        min_similarity: float = 0.0
    ) -> List[SearchResult]:
        """
        Search across all shards in parallel and merge results.

        Args:
            query_vector: Query vector
            k: Number of results to return
            min_similarity: Minimum similarity threshold

        Returns:
            List of SearchResult objects, sorted by similarity
        """
        shard_ids = self.shard_manager.get_all_shard_ids()

        # Search all shards in parallel
        all_results = []

        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            # Submit search tasks
            future_to_shard = {
                executor.submit(
                    self._search_shard,
                    shard_id,
                    query_vector,
                    k,
                    min_similarity
                ): shard_id
                for shard_id in shard_ids
            }

            # Collect results
            for future in as_completed(future_to_shard):
                shard_id = future_to_shard[future]
                try:
                    shard_results = future.result()
                    all_results.extend(shard_results)
                except Exception as e:
                    self._log(f"Error searching shard {shard_id}: {e}")

        # Merge and re-rank results
        merged_results = self._merge_results(all_results, k)

        self._log(
            f"Cross-shard search returned {len(merged_results)} results from {len(shard_ids)} shards")
        return merged_results

    def _search_shard(
        self,
        shard_id: int,
        query_vector: np.ndarray,
        k: int,
        min_similarity: float
    ) -> List[SearchResult]:
        """
        Search a single shard.

        Args:
            shard_id: Shard ID
            query_vector: Query vector
            k: Number of results
            min_similarity: Minimum similarity

        Returns:
            List of SearchResult objects
        """
        try:
            shard_index = self._get_shard_index(shard_id)
            # FaissIndex.search only takes (query, k)
            results = shard_index.search(query_vector, k=k)

            # Convert to SearchResult objects and filter by min_similarity
            search_results = []
            shard_meta = self._get_shard_metadata(shard_id)
            
            for episode_id, similarity in results:
                if similarity >= min_similarity:
                    search_results.append(SearchResult(
                        episode_id=episode_id,
                        similarity=similarity,
                        metadata=shard_meta.get(episode_id, {})
                    ))

            return search_results
        except Exception as e:
            self._log(f"Error in shard {shard_id}: {e}")
            return []

    def _merge_results(
        self,
        all_results: List[SearchResult],
        k: int
    ) -> List[SearchResult]:
        """
        Merge and re-rank results from multiple shards.

        Args:
            all_results: Combined results from all shards
            k: Number of top results to return

        Returns:
            Top-k results sorted by similarity
        """
        # Sort by similarity (descending)
        sorted_results = sorted(
            all_results,
            key=lambda x: x.similarity,
            reverse=True)

        # Return top-k
        return sorted_results[:k]

    def _update_shard_stats(
            self,
            shard_id: int,
            shard_index: FaissIndex) -> None:
        """Update shard statistics in the manager."""
        num_vectors = len(shard_index._ids)

        # Estimate disk size (rough approximation)
        disk_size = num_vectors * self.dimension * 4  # 4 bytes per float32
        self.shard_manager.update_shard_stats(shard_id, num_vectors, disk_size)

    def save_all_shards(self) -> None:
        """Save all loaded shards to disk."""
        for shard_id, shard_index in self._shard_cache.items():
            shard_path = self.shard_manager.get_shard_path(shard_id)
            shard_index.save(shard_path)
            
            # Save metadata
            if shard_id in self._metadata_cache:
                meta_path = shard_path.with_suffix(".meta.json")
                with open(meta_path, "w") as f:
                    json.dump(self._metadata_cache[shard_id], f)
            
            self._log(f"Saved shard {shard_id} to {shard_path}")

    def get_stats(self) -> Dict[str, Any]:
        """
        Get statistics for the sharded index.

        Returns:
            Dictionary with overall and per-shard statistics
        """
        shard_stats = self.shard_manager.get_shard_stats()

        total_vectors = sum(stats['num_vectors']
                            for stats in shard_stats.values())
        total_size_mb = sum(stats['disk_size_mb']
                            for stats in shard_stats.values())

        return {'num_shards': self.shard_manager.get_num_shards(),
                'total_vectors': total_vectors,
                'total_size_mb': total_size_mb,
                'avg_vectors_per_shard': total_vectors / max(1,
                                                             self.shard_manager.get_num_shards()),
                'shards': shard_stats,
                'needs_rebalancing': self.shard_manager.needs_rebalancing()}

    def close(self) -> None:
        """Save all shards and clear cache."""
        self.save_all_shards()
        self._shard_cache.clear()
        self._log("Closed sharded index")


__all__ = ["ShardedFaissIndex", "SearchResult"]
