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

from ask_handbook import select_quote
from mixed_search import DEFAULT_DB, retrieve


def candidates(database, as_of, part):
    equipment = set(re.findall(r"\bEQ-[0-9]{3}\b", part.upper()))
    if len(equipment) > 1:
        raise ValueError("Put each equipment ID in a separate question part.")
    if equipment:
        return retrieve(database, as_of, equipment=equipment.pop())
    return retrieve(database, as_of, query=part)


def answer(database, as_of, part, model):
    passages = candidates(database, as_of, part)
    return passages, select_quote(passages, part, model)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Check one source quote for each part of a question.")
    parser.add_argument("--db", type=Path, default=DEFAULT_DB)
    parser.add_argument("--as-of", type=date.fromisoformat, required=True)
    parser.add_argument("--model", default="qwen2.5:1.5b")
    parser.add_argument("--part", action="append", required=True)
    args = parser.parse_args()
    try:
        for number, part in enumerate(args.part, start=1):
            passages, result = answer(args.db, args.as_of, part, args.model)
            print(f"{number}. {part}")
            if result is None:
                print(f"   No quote verified from {len(passages)} candidates.")
            else:
                evidence, quote = result
                print(f"   {quote}\n   {evidence.citation()}")
    except (OSError, sqlite3.Error, ValueError) as error:
        raise SystemExit(str(error)) from None