import json
from pathlib import Path

SYSTEM = (
    "Classify this fictional library request. Use RETURN for giving back an item, "
    "RENEW for extending a loan, or RESERVE for booking equipment. "
    "Reply with exactly one label."
)

SPLITS = {
    "train": [
        ("Where can I drop off this book?", "RETURN"),
        ("I finished my book. Can I bring it back this evening?", "RETURN"),
        ("Which desk takes items back?", "RETURN"),
        ("Is there an after-hours drop box for my borrowed DVD?", "RETURN"),
        ("I am done with the atlas and want to give it back.", "RETURN"),
        ("How do I hand in the borrowed map?", "RETURN"),
        ("Can I keep this book for another week?", "RENEW"),
        ("Please extend my loan until next Thursday.", "RENEW"),
        ("The novel is due tomorrow and I need more time.", "RENEW"),
        ("Can you move the due date on my audiobook?", "RENEW"),
        ("How do I prolong my checkout for this item?", "RENEW"),
        ("I'd like the current borrowing period to last longer.", "RENEW"),
        ("Book the 3D printer for Saturday afternoon.", "RESERVE"),
        ("I need a time slot on the scanner.", "RESERVE"),
        ("Can I sign up for the sewing machine on Friday?", "RESERVE"),
        ("Please hold the projector for my presentation at noon.", "RESERVE"),
        ("I want to use the recording microphone at three.", "RESERVE"),
        ("Is there an open slot for the video camera tomorrow?", "RESERVE"),
    ],
    "valid": [
        ("Does the south desk accept borrowed board games back?", "RETURN"),
        ("Can my current checkout stay active through Monday?", "RENEW"),
        ("Could I schedule the button maker next Wednesday?", "RESERVE"),
    ],
    "test": [
        ("I found the library's cookbook in my bag. Where should I hand it over?", "RETURN"),
        ("Could my friend bring my finished magazine back for me?", "RETURN"),
        ("I am halfway through and my borrowing period ends today. May I have two more days?", "RENEW"),
        ("The checkout expires before my trip ends. Can I change the deadline?", "RENEW"),
        ("Please set aside the digitizing station for me after lunch.", "RESERVE"),
        ("I would like the last available time for the label printer.", "RESERVE"),
    ],
}

folder = Path("library_data")
folder.mkdir(exist_ok=False)
for name, examples in SPLITS.items():
    with (folder / f"{name}.jsonl").open("w", encoding="utf-8") as file:
        for request, label in examples:
            row = {"messages": [
                {"role": "system", "content": SYSTEM},
                {"role": "user", "content": request},
                {"role": "assistant", "content": label},
            ]}
            file.write(json.dumps(row) + "\n")
    print(f"{name}: {len(examples)} examples")