import os
import sys
from PIL import Image
import io
import torch

# Add the project root to sys.path to import src
sys.path.append(os.getcwd())

from src.api.main import ModelManager

def test_inference():
    print("Testing ModelManager initialization...")
    manager = ModelManager(onnx_path='best.onnx', pt_path='best.pt', benchmark_path='sawit_tbs.pt')
    print("ModelManager initialized successfully.")

    # Create a dummy image for testing
    img = Image.new('RGB', (640, 640), color = (73, 109, 137))
    
    print("\nTesting PyTorch inference (Native)...")
    detections, raw, ms = manager.run_pytorch_inference(img, 0.25, engine_type="pytorch")
    print(f"Detections: {len(detections)}, Inference: {ms:.2f}ms")
    
    print("\nTesting Benchmark inference (YOLOv8-Sawit)...")
    detections, raw, ms = manager.run_pytorch_inference(img, 0.25, engine_type="yolov8_sawit")
    print(f"Detections: {len(detections)}, Inference: {ms:.2f}ms")
    print(f"Benchmark Class Names: {manager.benchmark_class_names}")

    print("\nVerification Complete.")

if __name__ == "__main__":
    test_inference()