"""
index_shard_manager.py
----------------------
Manages FAISS index sharding with consistent hashing for distributed vector storage.

Provides automatic shard selection, load balancing, and rebalancing capabilities
for scaling to millions of vectors.

Author: HIEROMEM Team
"""

from __future__ import annotations
from typing import Any, Dict, List, Tuple
from pathlib import Path
import json
import mmh3
from threading import Lock
from dataclasses import dataclass, asdict
from datetime import datetime


@dataclass
class ShardMetadata:
    """Metadata for a single shard."""
    shard_id: int
    num_vectors: int
    disk_size_bytes: int
    created_at: str
    last_updated: str
    path: str


class IndexShardManager:
    """
    Manages FAISS index shards using consistent hashing.

    Features:
    - Consistent hashing with virtual nodes
    - Automatic shard creation and removal
    - Load-based rebalancing
    - Shard metadata tracking
    """

    def __init__(
        self,
        base_path: str | Path,
        num_shards: int = 4,
        virtual_nodes: int = 150,
        max_vectors_per_shard: int = 100000,
        debug: bool = False
    ):
        """
        Initialize the shard manager.

        Args:
            base_path: Base directory for shard storage
            num_shards: Initial number of shards
            virtual_nodes: Number of virtual nodes per shard for better distribution
            max_vectors_per_shard: Threshold for triggering rebalancing
            debug: Enable debug logging
        """
        self.base_path = Path(base_path)
        self.base_path.mkdir(parents=True, exist_ok=True)

        self.num_shards = num_shards
        self.virtual_nodes = virtual_nodes
        self.max_vectors_per_shard = max_vectors_per_shard
        self.debug = debug

        self._lock = Lock()
        self._shards: Dict[int, ShardMetadata] = {}
        self._hash_ring: List[Tuple[int, int]] = []  # (hash_value, shard_id)

        # Load or initialize shards
        self._load_or_initialize_shards()
        self._build_hash_ring()

        self._log(f"Initialized shard manager with {num_shards} shards")

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

    def _load_or_initialize_shards(self) -> None:
        """Load existing shards or initialize new ones."""
        metadata_file = self.base_path / "shard_metadata.json"

        if metadata_file.exists():
            # Load existing shards
            with open(metadata_file, 'r') as f:
                data = json.load(f)
                for shard_data in data['shards']:
                    shard = ShardMetadata(**shard_data)
                    self._shards[shard.shard_id] = shard
            self._log(f"Loaded {len(self._shards)} existing shards")
        else:
            # Initialize new shards
            for shard_id in range(self.num_shards):
                shard_path = self.base_path / f"shard_{shard_id}"
                shard_path.mkdir(exist_ok=True)

                shard = ShardMetadata(
                    shard_id=shard_id,
                    num_vectors=0,
                    disk_size_bytes=0,
                    created_at=datetime.now().isoformat(),
                    last_updated=datetime.now().isoformat(),
                    path=str(shard_path)
                )
                self._shards[shard_id] = shard

            self._save_metadata()
            self._log(f"Initialized {self.num_shards} new shards")

    def _save_metadata(self) -> None:
        """Save shard metadata to disk."""
        metadata_file = self.base_path / "shard_metadata.json"
        data = {
            'shards': [asdict(shard) for shard in self._shards.values()],
            'num_shards': len(self._shards),
            'virtual_nodes': self.virtual_nodes,
            'max_vectors_per_shard': self.max_vectors_per_shard
        }

        with open(metadata_file, 'w') as f:
            json.dump(data, f, indent=2)

    def _build_hash_ring(self) -> None:
        """Build the consistent hash ring with virtual nodes."""
        self._hash_ring = []

        for shard_id in self._shards.keys():
            for vnode in range(self.virtual_nodes):
                # Create virtual node key
                vnode_key = f"shard_{shard_id}_vnode_{vnode}"
                # Hash the virtual node key
                hash_value = mmh3.hash(vnode_key, signed=False)
                self._hash_ring.append((hash_value, shard_id))

        # Sort by hash value for binary search
        self._hash_ring.sort(key=lambda x: x[0])
        self._log(f"Built hash ring with {len(self._hash_ring)} virtual nodes")

    def get_shard_for_key(self, key: str) -> int:
        """
        Get the shard ID for a given key using consistent hashing.

        Args:
            key: Key to hash (e.g., episode_id)

        Returns:
            Shard ID
        """
        if not self._hash_ring:
            return 0

        # Hash the key
        key_hash = mmh3.hash(key, signed=False)

        # Binary search for the first hash value >= key_hash
        left, right = 0, len(self._hash_ring)
        while left < right:
            mid = (left + right) // 2
            if self._hash_ring[mid][0] < key_hash:
                left = mid + 1
            else:
                right = mid

        # Wrap around if necessary
        if left >= len(self._hash_ring):
            left = 0

        shard_id = self._hash_ring[left][1]
        return shard_id

    def add_shard(self) -> int:
        """
        Add a new shard to the cluster.

        Returns:
            New shard ID
        """
        with self._lock:
            # Find next available shard ID
            shard_id = max(self._shards.keys()) + 1 if self._shards else 0

            # Create shard directory
            shard_path = self.base_path / f"shard_{shard_id}"
            shard_path.mkdir(exist_ok=True)

            # Create metadata
            shard = ShardMetadata(
                shard_id=shard_id,
                num_vectors=0,
                disk_size_bytes=0,
                created_at=datetime.now().isoformat(),
                last_updated=datetime.now().isoformat(),
                path=str(shard_path)
            )

            self._shards[shard_id] = shard
            self._build_hash_ring()
            self._save_metadata()

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

    def remove_shard(self, shard_id: int) -> None:
        """
        Remove a shard from the cluster.

        WARNING: This does not migrate data. Ensure shard is empty first.

        Args:
            shard_id: Shard ID to remove
        """
        with self._lock:
            if shard_id not in self._shards:
                raise ValueError(f"Shard {shard_id} does not exist")

            shard = self._shards[shard_id]
            if shard.num_vectors > 0:
                raise ValueError(
                    f"Cannot remove shard {shard_id}: contains {shard.num_vectors} vectors")

            del self._shards[shard_id]
            self._build_hash_ring()
            self._save_metadata()

            self._log(f"Removed shard {shard_id}")

    def update_shard_stats(
            self,
            shard_id: int,
            num_vectors: int,
            disk_size_bytes: int) -> None:
        """
        Update statistics for a shard.

        Args:
            shard_id: Shard ID
            num_vectors: Current number of vectors
            disk_size_bytes: Current disk usage in bytes
        """
        with self._lock:
            if shard_id not in self._shards:
                raise ValueError(f"Shard {shard_id} does not exist")

            shard = self._shards[shard_id]
            shard.num_vectors = num_vectors
            shard.disk_size_bytes = disk_size_bytes
            shard.last_updated = datetime.now().isoformat()

            self._save_metadata()

    def get_shard_stats(self) -> Dict[int, Dict[str, Any]]:
        """
        Get statistics for all shards.

        Returns:
            Dictionary mapping shard_id to stats
        """
        return {
            shard_id: {
                'num_vectors': shard.num_vectors,
                'disk_size_bytes': shard.disk_size_bytes,
                'disk_size_mb': shard.disk_size_bytes / (1024 * 1024),
                'last_updated': shard.last_updated,
                'path': shard.path
            }
            for shard_id, shard in self._shards.items()
        }

    def needs_rebalancing(self) -> bool:
        """
        Check if any shard exceeds the size threshold.

        Returns:
            True if rebalancing is needed
        """
        for shard in self._shards.values():
            if shard.num_vectors > self.max_vectors_per_shard:
                self._log(
                    f"Shard {shard.shard_id} has {shard.num_vectors} vectors (limit: {self.max_vectors_per_shard})")
                return True
        return False

    def get_shard_path(self, shard_id: int) -> Path:
        """
        Get the filesystem path for a shard.

        Args:
            shard_id: Shard ID

        Returns:
            Path to shard directory
        """
        if shard_id not in self._shards:
            raise ValueError(f"Shard {shard_id} does not exist")
        return Path(self._shards[shard_id].path)

    def get_all_shard_ids(self) -> List[int]:
        """Get list of all shard IDs."""
        return list(self._shards.keys())

    def get_num_shards(self) -> int:
        """Get total number of shards."""
        return len(self._shards)


__all__ = ["IndexShardManager", "ShardMetadata"]
