from statistics import mean

from llama_index.core import VectorStoreIndex
from llama_index.core.embeddings import BaseEmbedding
from llama_index.core.evaluation import RetrieverEvaluator
from llama_index.core.schema import TextNode


class KeywordEmbedding(BaseEmbedding):
    def _vector_for_text(self, text: str) -> list[float]:
        text = text.lower()
        if "billing" in text and "release" in text:
            return [0.8, 0.6, 0.0]
        if "billing" in text or "7421" in text:
            return [1.0, 0.0, 0.0]
        if "release" in text:
            return [0.0, 1.0, 0.0]
        return [0.0, 0.0, 1.0]

    def _get_text_embedding(self, text: str) -> list[float]:
        return self._vector_for_text(text)

    def _get_query_embedding(self, query: str) -> list[float]:
        return self._vector_for_text(query)

    async def _aget_query_embedding(self, query: str) -> list[float]:
        return self._get_query_embedding(query)


nodes = [
    TextNode(
        id_="billing-escalation",
        text="Escalate stale billing answers for ticket 7421 to docs review.",
    ),
    TextNode(
        id_="release-freeze",
        text="Release freeze exceptions require approval from the release desk.",
    ),
    TextNode(
        id_="password-reset",
        text="Password reset requests require identity verification.",
    ),
]

index = VectorStoreIndex(nodes, embed_model=KeywordEmbedding())
retriever = index.as_retriever(similarity_top_k=2)

evaluator = RetrieverEvaluator.from_metric_names(
    ["hit_rate", "mrr"],
    retriever=retriever,
)

labeled_queries = [
    ("Where should billing answers for ticket 7421 be escalated?", ["billing-escalation"]),
    ("Who approves billing changes during the release freeze?", ["release-freeze"]),
]

results = []
for query, expected_ids in labeled_queries:
    result = evaluator.evaluate(query=query, expected_ids=expected_ids)
    results.append(result)
    print(f"query: {query}")
    print(f"expected_ids: {expected_ids}")
    print(f"retrieved_ids: {result.retrieved_ids}")
    print(f"hit_rate: {result.metric_vals_dict['hit_rate']}")
    print(f"mrr: {result.metric_vals_dict['mrr']}")
    print()

mean_hit_rate = mean(result.metric_vals_dict["hit_rate"] for result in results)
mean_mrr = mean(result.metric_vals_dict["mrr"] for result in results)
passed = mean_hit_rate == 1.0 and mean_mrr >= 0.75

print(f"mean_hit_rate: {mean_hit_rate}")
print(f"mean_mrr: {mean_mrr}")
print(f"status: {'PASS' if passed else 'FAIL'}")