from __future__ import annotations

import json
from pathlib import Path

from rag_pipeline import HandbookRAG, REFUSAL_TEXT


ROOT = Path(__file__).resolve().parent
GOLDEN_DATASET_PATH = ROOT / "golden_dataset.json"
RESULTS_PATH = ROOT / "test_results.json"


def load_golden_dataset() -> list[dict]:
    return json.loads(GOLDEN_DATASET_PATH.read_text(encoding="utf-8"))


def matches_keywords(answer: str, keywords: list[str]) -> bool:
    answer_lower = answer.lower()
    if not keywords:
        return True
    hits = 0
    for keyword in keywords:
        keyword_lower = keyword.lower()
        if keyword_lower in answer_lower:
            hits += 1
            continue

        words = [word for word in keyword_lower.split() if word]
        if words and all(word in answer_lower for word in words):
            hits += 1
    required = 1 if len(keywords) == 1 else min(2, len(keywords))
    return hits >= required


def is_refusal(answer: str) -> bool:
    answer_lower = answer.lower()
    refusal_markers = [
        "could not find enough support",
        "cannot answer from the handbook",
        "not enough support",
        "try rephrasing",
        "cannot be answered",
        "does not state a specific",
        "cannot answer this question",
        "does not provide",
    ]
    return any(marker in answer_lower for marker in refusal_markers) or answer.strip() == REFUSAL_TEXT


def main() -> None:
    payload = {
        "status": "completed",
        "results": [],
        "summary": {},
    }

    try:
        bot = HandbookRAG()
    except Exception as exc:
        payload["status"] = "not_run"
        payload["summary"] = {"reason": str(exc)}
        RESULTS_PATH.write_text(json.dumps(payload, indent=2), encoding="utf-8")
        print(json.dumps(payload, indent=2))
        return

    golden_dataset = load_golden_dataset()
    passed = 0

    for item in golden_dataset:
        try:
            answer, contexts = bot.answer(item["question"], [])
            retrieved_sources = [context.source for context in contexts]
            if item["answerable"]:
                source_ok = any(source in retrieved_sources for source in item["expected_sources"])
                keyword_ok = matches_keywords(answer, item["expected_keywords"])
                item_passed = source_ok and keyword_ok and not is_refusal(answer)
            else:
                item_passed = is_refusal(answer)
        except Exception as exc:
            answer = f"Error: {exc}"
            retrieved_sources = []
            item_passed = False

        if item_passed:
            passed += 1

        payload["results"].append(
            {
                "id": item["id"],
                "question": item["question"],
                "answerable": item["answerable"],
                "expected_sources": item["expected_sources"],
                "retrieved_sources": retrieved_sources,
                "answer": answer,
                "passed": item_passed,
            }
        )

    payload["summary"] = {
        "num_questions": len(golden_dataset),
        "num_passed": passed,
        "pass_rate": round(passed / len(golden_dataset), 4) if golden_dataset else 0.0,
    }
    RESULTS_PATH.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    print(json.dumps(payload["summary"], indent=2))


if __name__ == "__main__":
    main()