import os
import torch
from torchvision import transforms
from PIL import Image
import io

from model import get_model

class EndpointHandler:
    def __init__(self, path: str = ""):
        """
        Initialize the handler. Load the Faster R-CNN model.
        """
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model_weights_path = os.path.join(path, "model.pt") 
        
        # Load the model
        self.model = get_model(num_classes=4) 
        print(f"Loading weights from: {self.model_weights_path}")
        checkpoint = torch.load(self.model_weights_path, map_location=self.device)
        self.model.load_state_dict(checkpoint["model_state_dict"])
        self.model.to(self.device)
        self.model.eval()

        # Define image preprocessing
        self.transform = transforms.Compose([
            transforms.Resize((640, 640)),
            transforms.ToTensor(),
        ])

    def __call__(self, data):
        """
        Process the incoming request and return object detection predictions.
        """
        try:
            if "image" not in data:
                return [{"error": "No 'image' provided in request."}]

            image_bytes = data["image"].encode("latin1")
            image = Image.open(io.BytesIO(image_bytes)).convert("RGB")

            input_tensor = self.transform(image).unsqueeze(0).to(self.device)

            with torch.no_grad():
                outputs = self.model(input_tensor)

            boxes = outputs[0]["boxes"].cpu().tolist()
            labels = outputs[0]["labels"].cpu().tolist()
            scores = outputs[0]["scores"].cpu().tolist()

            threshold = 0.5
            predictions = [
                {"box": box, "label": label, "score": score}
                for box, label, score in zip(boxes, labels, scores)
                if score > threshold
            ]

            return [{"predictions": predictions}]
        except Exception as e:
            return [{"error": str(e)}]