import logging
import uuid
from typing import Any, Dict, List, Optional

from crewai.memory.storage.rag_storage import RAGStorage
from crewai.rag.embeddings.factory import build_embedder
from endee import Endee

from .utils import SUPPORTED_SPARSE_MODELS, get_model_vocab_size, get_sparse_encoder

_logger = logging.getLogger(__name__)


class EndeeVectorStore(RAGStorage):
    """
    High-performance vector store for CrewAI using Endee.

    Supports both dense-only and hybrid (dense + sparse) search modes.
    """

    TEST_STRING = "test"
    MAX_LENGTH_BYTES = 8192
    DEFAULT_SPARSE_MODEL = "splade_pp"
    app: Endee | None = None

    def __init__(
        self,
        type: str,
        allow_reset: bool = True,
        crew: Any = None,
        embedder_config: Optional[Any] = None,
        endee_index: Optional[Any] = None,
        text_key: str = "value",
        api_token: Optional[str] = None,
        space_type: str = "cosine",
        precision: Optional[str] = None,
        sparse_model_name: Optional[str] = None,
        sparse_dim: Optional[int] = None,
    ):
        """
        Initialize Endee vector store.

        Args:
            type: Unique index name
            allow_reset: Whether to allow index reset/deletion
            crew: CrewAI crew instance
            embedder_config: Configuration for dense embeddings
            endee_index: Pre-existing Endee index (optional)
            text_key: Metadata key for storing text content
            api_token: Endee API token (optional for local usage)
            space_type: Distance metric ('cosine', 'l2', 'ip')
            precision: Quantization level ('float32', 'float16', 'int16d', etc.)
            sparse_model_name: Model key for sparse embeddings (enables hybrid mode)
            sparse_dim: Vocabulary size for sparse vectors (auto-detected if None)
        """
        self.type = type
        self._api_token = api_token
        self.crew = crew
        self.embedder = build_embedder(embedder_config)
        self._space_type = space_type
        self._text_key = text_key
        self._vector_index = endee_index
        self._initialized = False
        self._dimension = None
        self._precision = precision

        # Hybrid Mode Logic: Enable if EITHER sparse_model_name OR sparse_dim is set
        self._hybrid = (sparse_model_name is not None) or (sparse_dim is not None)

        if self._hybrid:
            # Case 1: User provided sparse_dim ONLY -> Default to "splade_pp" model
            if sparse_model_name is None:
                self._sparse_model_name = self.DEFAULT_SPARSE_MODEL
                _logger.info(
                    f"Hybrid mode triggered by 'sparse_dim'. Defaulting model to '{self._sparse_model_name}'."
                )
            else:
                self._sparse_model_name = sparse_model_name

            # Case 2: User provided model ONLY (or we just defaulted it) -> Auto-detect Dim
            if sparse_dim is None:
                self._sparse_dim = get_model_vocab_size(self._sparse_model_name)
                _logger.info(
                    f"Auto-detected sparse_dim={self._sparse_dim} for model '{self._sparse_model_name}'"
                )
            else:
                self._sparse_dim = sparse_dim
                _logger.info(
                    f"Using provided sparse_dim={self._sparse_dim} for model '{self._sparse_model_name}'"
                )

            # Validate model support
            if self._sparse_model_name not in SUPPORTED_SPARSE_MODELS:
                _logger.warning(
                    f"Model '{self._sparse_model_name}' not in SUPPORTED_SPARSE_MODELS. "
                    f"Supported: {list(SUPPORTED_SPARSE_MODELS.keys())}"
                )

            # Initialize sparse encoder
            try:
                self._sparse_encoder = get_sparse_encoder(
                    self._sparse_model_name, use_fastembed=True
                )
                _logger.info(
                    f"Hybrid mode initialized with model '{self._sparse_model_name}'"
                )
            except Exception as e:
                _logger.error(f"Failed to initialize sparse encoder: {e}")
                raise
        else:
            # Dense-only mode
            self._sparse_model_name = None
            self._sparse_dim = None
            self._sparse_encoder = None
            _logger.info("Dense-only mode enabled")

        super().__init__(type, allow_reset, embedder_config, crew)

    def ensure_index(self):
        """
        Public method to explicitly initialize index with clear output.

        Returns:
            The initialized Endee index
        """
        mode = "hybrid" if self._hybrid else "dense-only"
        print(f"\n[Endee] Ensuring index '{self.type}' exists ({mode} mode)...")
        self._ensure_initialized()

        if self._hybrid:
            print(
                f"[Endee] Index '{self.type}' ready with sparse_dim={self._sparse_dim}\n"
            )
        else:
            print(f"[Endee] Index '{self.type}' ready.\n")

        return self._vector_index

    def _ensure_initialized(self):
        """Ensure the Endee app and index are initialized before any operation."""
        if not self._initialized:
            self._initialize_app()
            self._initialized = True

    def _initialize_app(self):
        """Initialize or retrieve the Endee index."""
        # Determine dense embedding dimension
        if self._dimension is None:
            self._dimension = len(self.embedder([self.TEST_STRING])[0])
            _logger.info(f"Dense embedding dimension: {self._dimension}")

        # Initialize Endee client
        if self.app is None:
            if self._api_token:
                self.app = Endee(token=self._api_token)
                _logger.info("Connected to Endee with API token")
            else:
                self.app = Endee()
                _logger.info("Using Endee in local mode")

        # Try to retrieve existing index
        try:
            index = self.app.get_index(name=self.type)
            self._vector_index = index
            _logger.info(f"Retrieved existing index: {self.type}")
            return index
        except Exception as e:
            _logger.debug(f"Index '{self.type}' does not exist: {e}")

        # Create new index
        _logger.info(
            f"Creating new index: {self.type} "
            f"(dimension={self._dimension}, space={self._space_type}, "
            f"precision={self._precision}, hybrid={self._hybrid})"
        )

        try:
            if self._hybrid:
                self.app.create_index(
                    name=self.type,
                    dimension=self._dimension,
                    space_type=self._space_type,
                    precision=self._precision,
                    sparse_dim=self._sparse_dim,
                )
                _logger.info(f"Created hybrid index with sparse_dim={self._sparse_dim}")
            else:
                self.app.create_index(
                    name=self.type,
                    dimension=self._dimension,
                    space_type=self._space_type,
                    precision=self._precision,
                )
                _logger.info("Created dense-only index")

            self._vector_index = self.app.get_index(name=self.type)
            return self._vector_index

        except Exception as create_error:
            _logger.error(f"Error creating index: {create_error}")
            raise

    def _compute_sparse_vectors(self, texts: List[str]) -> tuple:
        """
        Compute sparse vectors for a list of texts.

        Args:
            texts: List of text strings

        Returns:
            Tuple of (indices, values) where each is a list of lists
        """
        try:
            if self._sparse_encoder is None:
                raise ValueError(
                    "Sparse encoder not initialized. "
                    "Please provide sparse_model_name when creating the store."
                )

            indices, values = self._sparse_encoder(texts)

            # Validate indices are within bounds
            for idx_list in indices:
                if idx_list and max(idx_list) >= self._sparse_dim:
                    raise ValueError(
                        f"Sparse indices out of bounds [0, {self._sparse_dim}). "
                        f"Max index found: {max(idx_list)}. "
                        f"This may indicate sparse_dim mismatch with the model."
                    )

            return indices, values

        except Exception as e:
            _logger.error(f"Error computing sparse vectors: {e}")
            raise

    def save(self, value: str, metadata: Dict[str, Any]) -> None:
        """
        Insert a new vector into the Endee index.

        Args:
            value: Text content to store
            metadata: Associated metadata dictionary
        """
        self._ensure_initialized()

        # Normalize text to fit byte limits
        value = self._normalize_text(value)

        # Compute dense embedding
        embedding = self.embedder([value])[0]

        # Add text to metadata
        metadata[self._text_key] = value

        # Compute sparse vectors if in hybrid mode
        indices, values = None, None
        if self._hybrid:
            indices, values = self._compute_sparse_vectors([value])

        # Prepare filter dictionary, We copy all strings, ints, and floats from metadata into the filter.
        filter_dict = {
            k: v for k, v in metadata.items() if isinstance(v, (str, int, float))
        }

        # Prepare event for upsert
        try:
            event = {
                "id": uuid.uuid4().hex,
                "vector": embedding,
                "meta": metadata,
                "filter": filter_dict,
            }

            if self._hybrid:
                event["sparse_indices"] = indices[0]
                event["sparse_values"] = values[0]

            self._vector_index.upsert([event])
            _logger.debug(f"Saved document with id={event['id']}")

        except Exception as e:
            _logger.error(f"Error saving data to Endee: {str(e)}")
            raise

    def search(
        self,
        query: str,
        limit: int = 3,
        filter: Optional[List[dict[str, Any]]] = None,
        score_threshold: float = 0,
    ) -> list[dict]:
        """
        Search for similar vectors in the Endee index.

        Args:
            query: Search query text
            limit: Maximum number of results to return
            filter: Optional metadata filters
            score_threshold: Minimum similarity score (0-1)

        Returns:
            List of search results with content, metadata, and scores
        """
        self._ensure_initialized()

        # Normalize query
        query = self._normalize_text(query)

        # Compute dense embedding
        embedding = self.embedder([query])[0]

        # Compute sparse vectors if in hybrid mode
        indices, values = None, None
        if self._hybrid:
            indices, values = self._compute_sparse_vectors([query])

        # Prepare query parameters
        try:
            kwargs = {
                "vector": embedding,
                "top_k": limit,
                "filter": filter or [],
                "include_vectors": True,
            }

            if self._hybrid:
                kwargs.update(
                    {
                        "sparse_indices": indices[0],
                        "sparse_values": values[0],
                    }
                )

            results = self._vector_index.query(**kwargs)
            _logger.debug(f"Search returned {len(results)} results")

        except Exception as e:
            _logger.error(f"Error querying Endee: {e}")
            return []

        # Process results
        final_results = []
        for result in results:
            try:
                score = result.get("similarity", 0)

                # Apply score threshold
                if score >= score_threshold:
                    content = result.get("meta", {}).get(self._text_key, "")
                    final_results.append(
                        {
                            "id": result.get("id", ""),
                            "metadata": result.get("meta", {}),
                            "context": content,
                            "content": content,
                            "score": score,
                        }
                    )
            except Exception as e:
                _logger.warning(
                    f"Error processing search result: {e}, result: {result}"
                )
                continue

        return final_results

    def reset(self) -> None:
        """
        Delete the entire Endee index and reset internal state.

        Warning: This permanently deletes all data in the index.
        """
        if self.app is None:
            if self._api_token:
                self.app = Endee(token=self._api_token)
            else:
                self.app = Endee()

        try:
            self.app.delete_index(self.type)
            _logger.info(f"Reset/deleted index: {self.type}")
        except Exception as e:
            _logger.info(
                f"Index {self.type} did not exist or could not be deleted: {e}"
            )

        # Reset internal state
        self._vector_index = None
        self._initialized = False

    def _normalize_text(self, text: str) -> str:
        """
        Normalize text to fit within the maximum byte length.

        Args:
            text: Input text

        Returns:
            Normalized text that fits within MAX_LENGTH_BYTES
        """
        encoded = text.encode("utf-8")
        if len(encoded) > self.MAX_LENGTH_BYTES:
            encoded = encoded[: self.MAX_LENGTH_BYTES]
            _logger.debug(f"Text truncated to {self.MAX_LENGTH_BYTES} bytes")
        return encoded.decode("utf-8", errors="ignore")
