from dataclasses import dataclass
import numpy as np
from src.labdaps.config import TOP_K_CANDIDATES, TOP_K_FINAL
from src.labdaps.retrieval.vector_store import query_store
from src.labdaps.ingestion.embedder import Embedder


@dataclass
class RetrievedChunk:
    text: str
    source_file: str
    page_number: int
    score: float


def _cosine_similarity(a: list[float], b: list[float]) -> float:
    a, b = np.array(a), np.array(b)
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-9))


def _mmr(
    query_embedding: list[float],
    candidate_embeddings: list[list[float]],
    candidates: list[dict],
    top_k: int,
    lambda_param: float = 0.6,
) -> list[int]:
    selected = []
    remaining = list(range(len(candidates)))
    query_sims = [_cosine_similarity(query_embedding, e) for e in candidate_embeddings]

    while len(selected) < top_k and remaining:
        if not selected:
            best = max(remaining, key=lambda i: query_sims[i])
        else:
            selected_embeddings = [candidate_embeddings[i] for i in selected]
            best = max(
                remaining,
                key=lambda i: lambda_param * query_sims[i]
                - (1 - lambda_param) * max(
                    _cosine_similarity(candidate_embeddings[i], se)
                    for se in selected_embeddings
                ),
            )
        selected.append(best)
        remaining.remove(best)

    return selected


def retrieve(query: str, embedder: Embedder) -> list[RetrievedChunk]:
    query_embedding = embedder.embed_query(query)
    documents, metadatas, distances = query_store(query_embedding, TOP_K_CANDIDATES)

    if not documents:
        return []

    candidate_embeddings = embedder.embed_passages(documents)
    selected_indices = _mmr(query_embedding, candidate_embeddings, metadatas, TOP_K_FINAL)

    return [
        RetrievedChunk(
            text=documents[idx],
            source_file=metadatas[idx]["source_file"],
            page_number=metadatas[idx]["page_number"],
            score=1.0 - distances[idx],
        )
        for idx in selected_indices
    ]