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

# Import your Faster R-CNN model definition
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")  # Adjust for your file name
        
        # Load the model
        self.model = get_model(num_classes=4)  # Modify for your num_classes
        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)),  # Adjust size to match your training setup
            transforms.ToTensor(),
        ])

    def __call__(self, data):
        """
        Process the incoming request and return object detection predictions.
        """
        try:
            # Expect input data to include a Base64-encoded image
            if "image" not in data:
                return [{"error": "No 'image' provided in request."}]

            # Convert Base64-encoded image to bytes
            image_bytes = data["image"].encode("latin1")
            image = Image.open(io.BytesIO(image_bytes)).convert("RGB")

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

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

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

            # Confidence threshold
            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)}]