# Copyright (C) 2026 GoodData Corporation

"""In-memory vector store for rule embeddings."""

from __future__ import annotations

import logging
from dataclasses import dataclass

import numpy as np
from numpy.typing import NDArray

from aida_mcp.tools.rules.embeddings.chunker import RuleChunk

logger = logging.getLogger(__name__)


@dataclass
class SearchResult:
    """Result from vector similarity search."""

    chunk: RuleChunk
    score: float  # Cosine similarity (0-1, higher is more similar)


class VectorStore:
    """
    Simple in-memory vector store using numpy.

    Stores embeddings and performs cosine similarity search.
    Designed to be recreated on each server startup.
    """

    def __init__(self) -> None:
        """Initialize empty vector store."""
        self._chunks: list[RuleChunk] = []
        self._embeddings: NDArray[np.float32] | None = None

    def add(self, chunks: list[RuleChunk], embeddings: NDArray[np.float32]) -> None:
        """
        Add chunks with their embeddings to the store.

        Args:
            chunks: List of rule chunks
            embeddings: Numpy array of shape (len(chunks), dimension)
        """
        if len(chunks) != len(embeddings):
            msg = f"Chunks ({len(chunks)}) and embeddings ({len(embeddings)}) must have same length"
            raise ValueError(msg)

        self._chunks.extend(chunks)

        if self._embeddings is None:
            self._embeddings = embeddings
        else:
            self._embeddings = np.vstack([self._embeddings, embeddings])

        logger.debug(f"Added {len(chunks)} chunks to vector store (total: {len(self._chunks)})")

    def search(
        self,
        query_embedding: NDArray[np.float32],
        top_k: int = 10,
        threshold: float = 0.0,
    ) -> list[SearchResult]:
        """
        Search for most similar chunks.

        Args:
            query_embedding: Query embedding vector (dimension,)
            top_k: Maximum number of results
            threshold: Minimum similarity score (0-1)

        Returns:
            List of SearchResult sorted by similarity (highest first)
        """
        if self._embeddings is None or len(self._chunks) == 0:
            return []

        # Compute cosine similarity (embeddings are already normalized)
        # dot product of normalized vectors = cosine similarity
        similarities = np.dot(self._embeddings, query_embedding)

        # Get top-k indices
        top_indices = np.argsort(similarities)[::-1][:top_k]

        results: list[SearchResult] = []
        for idx in top_indices:
            i = int(idx)  # Convert numpy int to Python int
            score = float(similarities[i])
            if score >= threshold:
                results.append(SearchResult(chunk=self._chunks[i], score=score))

        return results

    def search_by_rule_path(
        self,
        query_embedding: NDArray[np.float32],
        top_k: int = 10,
        threshold: float = 0.0,
    ) -> dict[str, float]:
        """
        Search and aggregate scores by rule path.

        For rules with multiple chunks, returns the maximum similarity score.

        Args:
            query_embedding: Query embedding vector
            top_k: Maximum number of chunks to consider
            threshold: Minimum similarity score

        Returns:
            Dict mapping rule_path to best similarity score
        """
        results = self.search(query_embedding, top_k=top_k * 3, threshold=threshold)

        # Aggregate by rule path (take max score per rule)
        scores_by_path: dict[str, float] = {}
        for result in results:
            path = result.chunk.rule_path
            if path not in scores_by_path or result.score > scores_by_path[path]:
                scores_by_path[path] = result.score

        # Sort by score and limit to top_k
        sorted_paths = sorted(scores_by_path.items(), key=lambda x: x[1], reverse=True)
        return dict(sorted_paths[:top_k])

    def search_with_chunks(
        self,
        query_embedding: NDArray[np.float32],
        top_k: int = 10,
        chunks_per_rule: int = 3,
        threshold: float = 0.0,
    ) -> tuple[dict[str, float], dict[str, list[SearchResult]]]:
        """
        Search and return both scores and matching chunks grouped by rule path.

        For large documents that are chunked, this allows returning only the
        most relevant chunks instead of the entire document.

        Args:
            query_embedding: Query embedding vector
            top_k: Maximum number of rules to return
            chunks_per_rule: Maximum chunks to return per rule
            threshold: Minimum similarity score

        Returns:
            Tuple of:
            - Dict mapping rule_path to best similarity score
            - Dict mapping rule_path to list of top matching SearchResults
        """
        # Get more results to ensure we have enough chunks per rule
        results = self.search(query_embedding, top_k=top_k * chunks_per_rule * 2, threshold=threshold)

        # Group by rule path
        scores_by_path: dict[str, float] = {}
        chunks_by_path: dict[str, list[SearchResult]] = {}

        for result in results:
            path = result.chunk.rule_path

            # Track best score per rule
            if path not in scores_by_path or result.score > scores_by_path[path]:
                scores_by_path[path] = result.score

            # Collect chunks per rule (up to limit, skip description-only chunks)
            if path not in chunks_by_path:
                chunks_by_path[path] = []
            if len(chunks_by_path[path]) < chunks_per_rule and not result.chunk.is_description:
                chunks_by_path[path].append(result)

        # Sort by score and limit to top_k rules
        sorted_paths = sorted(scores_by_path.items(), key=lambda x: x[1], reverse=True)[:top_k]
        top_paths = {path for path, _ in sorted_paths}

        # Filter to only top_k rules
        filtered_scores = {path: score for path, score in scores_by_path.items() if path in top_paths}
        filtered_chunks = {path: chunks for path, chunks in chunks_by_path.items() if path in top_paths}

        return filtered_scores, filtered_chunks

    @property
    def size(self) -> int:
        """Number of chunks in the store."""
        return len(self._chunks)

    def clear(self) -> None:
        """Clear all data from the store."""
        self._chunks = []
        self._embeddings = None
