from fastapi import FastAPI, UploadFile
from ultralytics import YOLO
import io
from PIL import Image
import numpy as np
import os
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
import requests
import supervision as sv

# Global variable to store model instances
MODEL_CACHE = {}

def init_model(model_id: str):
    # Return cached model if it exists
    if model_id in MODEL_CACHE:
        return MODEL_CACHE[model_id]
    
    # Define models
    MODEL_OPTIONS = {
        "YOLOv11-Nano": "medieval-yolov11n.pt",
        "YOLOv11-Small": "medieval-yolov11s.pt",
        "YOLOv11-Medium": "medieval-yolov11m.pt",
        "YOLOv11-Large": "medieval-yolov11l.pt",
        "YOLOv11-XLarge": "medieval-yolov11x.pt"
    }
    
    if model_id in MODEL_OPTIONS:
        os.makedirs("models", exist_ok=True)
        model_path = hf_hub_download(
            repo_id="biglam/medieval-manuscript-yolov11",
            filename=MODEL_OPTIONS[model_id],
            cache_dir="models"  # Specify cache directory
        )
        model = YOLO(model_path)
        MODEL_CACHE[model_id] = model
        return model
    else:
        raise ValueError(f"Model {model_id} not found")

app = FastAPI()

# Initialize default model at startup
@app.on_event("startup")
async def startup_event():
    init_model("YOLOv11-XLarge")  # Initialize default model

@app.post("/predict")
async def predict(image: UploadFile,
                  model_id: str = "YOLOv11-XLarge",
                  conf: float = 0.25,
                  iou: float = 0.7
                  ):
    # Get model from cache or initialize it
    model = init_model(model_id)
    
    # Download and open image from URL
    image = Image.open(image.file)
    
    # Run inference with the PIL Image
    results = model.predict(source=image, conf=conf, iou=iou)
    
    # Extract detection results
    result = results[0]
    # print(result)
    detections = []
    
    for box in result.boxes:
        detection = {
            "class": result.names[int(box.cls[0])],
            "confidence": float(box.conf[0]),
            "bbox": box.xyxy[0].tolist()  # Convert bbox tensor to list
        }
        detections.append(detection)
    print(detections)

    return {"detections": detections}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)