import argparse
import sys
from ultralytics import YOLO

def main(image_path: str, model_path: str = 'aerialEye.pt', conf: float = 0.25):
    """
    Run inference on an image using the trained YOLO model.
    """
    import os
    from download_model import is_lfs_pointer, download_file

    print(f"Loading model from {model_path}...")
    try:
        if model_path in ['aerialEye.pt', 'best.pt', 'aerialEye.onnx', 'best.onnx'] and (not os.path.exists(model_path) or is_lfs_pointer(model_path)):
            print(f"{model_path} is missing or is an LFS pointer. Downloading from Hugging Face...")
            download_file(model_path)
        model = YOLO(model_path)
    except Exception as e:
        print(f"Error loading model: {e}")
        sys.exit(1)
        
    print(f"Running inference on {image_path}...")
    results = model.predict(source=image_path, conf=conf, save=True)
    
    print("\nInference Complete!")
    print(f"Results saved to: {results[0].save_dir}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run YOLO Patrol Model Inference")
    parser.add_argument("--image", type=str, required=True, help="Path to the input image or video")
    parser.add_argument("--model", type=str, default="aerialEye.pt", help="Path to the YOLOv8 model")
    parser.add_argument("--conf", type=float, default=0.25, help="Confidence threshold")
    
    args = parser.parse_args()
    main(args.image, args.model, args.conf)