import argparse
import base64
import json
from pathlib import Path

from semantic_notes import api

FIELDS = ("time", "room", "signup")
EXPECTED = {"time": "14:30", "room": "204", "signup": "front desk"}
SCHEMA = {
    "type": "object",
    "properties": {field: {"type": ["string", "null"]} for field in FIELDS},
    "required": list(FIELDS),
    "additionalProperties": False,
}

parser = argparse.ArgumentParser(description="Check a model's extraction on a workshop notice.")
parser.add_argument("image", type=Path)
parser.add_argument("--model", default="gemma3:4b")
parser.add_argument("--hidden", choices=FIELDS, help="Field covered in your edited image")
args = parser.parse_args()

try:
    image = base64.b64encode(args.image.read_bytes()).decode("ascii")
    response = api("/api/chat", {
        "model": args.model,
        "messages": [{
            "role": "user",
            "content": (
                "Transcribe three fields from the notice. Set time to the exact time, "
                "room to the exact room number, and signup to the place name only, "
                "with no leading preposition or article. "
                "Use null for a field you cannot read. Return the fields as JSON."
            ),
            "images": [image],
        }],
        "format": SCHEMA,
        "options": {"temperature": 0, "num_ctx": 4096, "num_predict": 128},
        "stream": False,
    })
    result = json.loads(response["message"]["content"])
    if not isinstance(result, dict) or set(result) != set(FIELDS):
        raise ValueError("The model did not return all three fields.")
except (OSError, KeyError, TypeError, ValueError) as error:
    raise SystemExit(f"Could not score the image: {error}") from None

correct = 0
for field in FIELDS:
    expected = None if field == args.hidden else EXPECTED[field]
    actual = result[field]
    if isinstance(actual, str):
        actual = " ".join(actual.casefold().strip(" .").split())
    passed = actual == expected
    correct += passed
    print(f"{'PASS' if passed else 'MISS'} | {field}: {actual!r} | expected {expected!r}")
print(f"{correct}/{len(FIELDS)} matched")