import argparse

from semantic_notes import api, cosine

PASSAGES = {
    "renewal-zh": "如果没有其他读者预约，你可以把借来的图书续借十四天。",
    "hours-es": "La sala de lectura cierra a las seis de la tarde los sábados.",
    "scanner-zh": "设备编号 EQ-314 对应平板扫描仪。请在媒体服务台预约。",
    "tripod-en": "Borrow a camera tripod from the media desk. You can keep it for two days.",
    "returns-en": "Return books to the front desk or the after-hours book drop.",
    "workshop-es": "El taller de escritura empieza a las dos en la sala principal.",
}
CASES = [
    ("How long can I extend a book loan if nobody is waiting?", "renewal-zh"),
    ("周六阅览室几点关门？", "hours-es"),
    ("¿Qué equipo corresponde a EQ-314?", "scanner-zh"),
    ("¿Dónde puedo pedir prestado un trípode para la cámara?", "tripod-en"),
    ("晚上可以把书还到哪里？", "returns-en"),
    ("What time does the writing workshop begin?", "workshop-es"),
]


def embeddings(model, texts):
    values = api("/api/embed", {"model": model, "input": texts, "truncate": False})["embeddings"]
    if len(values) != len(texts):
        raise ValueError("Ollama returned the wrong number of embeddings.")
    return values


def evaluate(model):
    texts = list(PASSAGES.values())
    queries = [question for question, _ in CASES]
    if model == "embeddinggemma":
        texts = ["title: none | text: " + text for text in texts]
        queries = ["task: search result | query: " + query for query in queries]
    documents = dict(zip(PASSAGES, embeddings(model, texts)))
    questions = embeddings(model, queries)
    first = within_three = 0
    print(model)
    for (question, expected), vector in zip(CASES, questions):
        ranked = sorted(documents, key=lambda key: cosine(documents[key], vector), reverse=True)
        first += ranked[0] == expected
        within_three += expected in ranked[:3]
        print(f"  {question}\n    expected {expected}; top three: {', '.join(ranked[:3])}")
    print(f"  First result: {first}/{len(CASES)} | Top three: {within_three}/{len(CASES)}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Evaluate cross-language retrieval on six fictional notes.")
    parser.add_argument("model", choices=("all-minilm", "embeddinggemma", "both"), default="both", nargs="?")
    args = parser.parse_args()
    try:
        for name in ("all-minilm", "embeddinggemma") if args.model == "both" else (args.model,):
            evaluate(name)
    except (OSError, KeyError, ValueError) as error:
        raise SystemExit(str(error)) from None