# Copyright (C) 2026 GoodData Corporation

"""Embedding-based scoring for rules retrieval."""

from __future__ import annotations

import logging

from aida_mcp.tools.rules.embeddings.chunker import RuleChunk
from aida_mcp.tools.rules.embeddings.model import get_embedding_model
from aida_mcp.tools.rules.embeddings.store import VectorStore
from aida_mcp.tools.rules.store import Rule

logger = logging.getLogger(__name__)


class EmbeddingScorer:
    """Embedding-based semantic similarity scorer."""

    DEFAULT_CHUNKS_PER_RULE = 3

    def __init__(self, vector_store: VectorStore, chunks_per_rule: int = DEFAULT_CHUNKS_PER_RULE):
        self._vector_store = vector_store
        self._chunks_per_rule = chunks_per_rule

    def score_rules(self, query: str, rules: list[Rule]) -> tuple[dict[str, float], dict[str, list[RuleChunk]]]:
        empty_chunks: dict[str, list[RuleChunk]] = {}

        if self._vector_store.size == 0:
            logger.warning("Vector store is empty, returning zero scores")
            return {r.path: 0.0 for r in rules}, empty_chunks

        model = get_embedding_model()
        query_embedding = model.embed_query(query)

        scores, chunks_by_path = self._vector_store.search_with_chunks(
            query_embedding,
            top_k=len(rules),
            chunks_per_rule=self._chunks_per_rule,
            threshold=0.0,
        )

        matched_chunks: dict[str, list[RuleChunk]] = {
            path: [result.chunk for result in results] for path, results in chunks_by_path.items()
        }

        return {r.path: scores.get(r.path, 0.0) for r in rules}, matched_chunks
