from __future__ import annotations

import argparse
from pathlib import Path

from PIL import Image

from model_service import get_model_config, get_model_service


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Run local inference with the configured model.")
    parser.add_argument("--image", type=Path, default=Path("person.jpg"), help="Input image path.")
    parser.add_argument(
        "--model",
        type=str,
        default=None,
        help="Optional model name override. Defaults to SECUREML_MODEL or the project default.",
    )
    return parser


def main() -> None:
    args = build_parser().parse_args()
    if not args.image.exists():
        raise SystemExit(f"Image not found: {args.image}")

    config = get_model_config(args.model)
    service = get_model_service(args.model)
    image = Image.open(args.image).convert("RGB")
    result = service.predict_image(image)

    print(f"[INFO] device={service.device}")
    print(f"[INFO] model_name={config.name}")
    print(f"[INFO] model_backend={config.backend}")
    print(f"[INFO] model_path={config.model_path}")
    print(f"[INFO] image={args.image}")
    print("========== RESULT ==========")
    print(f"prediction: {result['prediction_index']} ({result['label']})")
    for label, prob in result["probabilities"].items():
        print(f"P({label}) = {prob:.6f}")
    print("============================")


if __name__ == "__main__":
    main()