import argparse
from datetime import date
from pathlib import Path
import sqlite3

from ask_handbook import select_quote
from mixed_answer import candidates
from mixed_search import DEFAULT_DB, retrieve

# Date, question, expected source record, and a phrase a correct quote should contain.
CASES = (
    ("2026-07-15", "What is the loan limit for EQ-314?",
     ("equipment-policy-2026-06-01.pdf", 1, "table", 2), "30 minutes"),
    ("2026-11-14", "What is the loan limit for EQ-314?",
     ("equipment-policy-2026-11-01.pdf", 1, "table", 2), "45 minutes"),
    ("2026-11-14", "When does the desk accept maps for the community map exchange?",
     ("library-handbook.pdf", 4, "ocr", None), "12:00 to 15:00"),
    ("2026-11-14", "How many days does a standard book loan last?",
     ("library-handbook.pdf", 1, "prose", None), "21 days"),
    ("2026-11-14", "What is the loan limit for EQ-073?", None, None),
    ("2026-11-14", "Who founded the library?", None, None),
)
COMPOUND_QUESTION = (
    "What is the loan limit for EQ-314, when does the desk accept maps for the community map exchange, "
    "and how many days does a standard book loan last?"
)


def location(item):
    return item.source, item.page, item.kind, item.row_number


def evaluate(database, model):
    retrieval_passes = answer_passes = 0
    for when, question, expected, phrase in CASES:
        items = candidates(database, date.fromisoformat(when), question)
        retrieval_ok = not items if expected is None else any(location(item) == expected for item in items)
        retrieval_passes += retrieval_ok
        print(f"{when} | {question}")
        print(f"  Retrieval: {'pass' if retrieval_ok else 'miss'} | {len(items)} candidates")
        if model:
            result = select_quote(items, question, model)
            if result is None:
                matched = expected is None
                print(f"  Answer: {'pass' if matched else 'miss'} | no quote")
            else:
                item, quote = result
                matched = expected is not None and location(item) == expected and phrase in quote
                print(f"  Answer: {'pass' if matched else 'miss'} | {quote}")
                print(f"  {item.citation()}")
            answer_passes += matched
    print(f"Retrieval checks: {retrieval_passes}/{len(CASES)}")
    if model:
        print(f"Answer checks: {answer_passes}/{len(CASES)}")


def compare_coverage(database):
    parts = CASES[1:4]
    as_of = date.fromisoformat(parts[0][0])
    combined = {location(item) for item in retrieve(database, as_of, query=COMPOUND_QUESTION)}
    print("Coverage by claim | one search / separate searches")
    for _, question, expected, _ in parts:
        separate = {location(item) for item in candidates(database, as_of, question)}
        before = "found" if expected in combined else "missed"
        after = "found" if expected in separate else "missed"
        print(f"  {before:6} / {after:6} | {question}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Check retrieval and answers on six mixed-document cases.")
    parser.add_argument("--db", type=Path, default=DEFAULT_DB)
    parser.add_argument("--model", help="Run the answer checks using this local Ollama model")
    parser.add_argument("--coverage", action="store_true", help="Compare retrieval for the compound question")
    args = parser.parse_args()
    if args.model and args.coverage:
        parser.error("Choose --coverage or --model for this run.")
    try:
        if args.coverage:
            compare_coverage(args.db)
        else:
            evaluate(args.db, args.model)
    except (OSError, sqlite3.Error, ValueError) as error:
        raise SystemExit(str(error)) from None