"""Small text-only inference example for the self-contained FxAI checkpoint."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

# Running the example inside an upload folder should not create *.pyc files.
sys.dont_write_bytecode = True

from futurix_decision import DecisionInput, load_model


def parse_option(value: str) -> tuple[str, str]:
    name, separator, description = value.partition("=")
    if not separator or not name.strip() or not description.strip():
        raise argparse.ArgumentTypeError("options must use NAME=DESCRIPTION")
    return name.strip(), description.strip()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument(
        "--state",
        default="My integration keeps failing. Please help.",
        help="Text to classify.",
    )
    parser.add_argument(
        "--instructions",
        default="Which team should handle this request?",
        help="Question shown to the decision model.",
    )
    parser.add_argument(
        "--option",
        action="append",
        type=parse_option,
        metavar="NAME=DESCRIPTION",
        help="Candidate decision option; repeat for multiple options.",
    )
    parser.add_argument("--max-image-pixels", type=int, default=512 * 512)
    args = parser.parse_args()
    options = args.option or [
        ("billing", "Charges and refunds"),
        ("support", "Technical integration errors"),
        ("sales", "Questions about purchasing the product"),
    ]
    if len({name for name, _ in options}) != len(options):
        parser.error("option names must be unique")

    model = load_model(
        args.checkpoint,
        family="fxai",
        device=args.device,
        max_image_pixels=args.max_image_pixels,
    )
    row = DecisionInput(
        state=args.state,
        question=args.instructions,
        options=[name for name, _ in options],
        candidate_descriptions=[description for _, description in options],
    )
    print(json.dumps(model.decide(row).to_dict(), indent=2))


if __name__ == "__main__":
    main()