from typing import Dict, List, Any
import torch
import numpy as np
from PIL import Image
from io import BytesIO
import base64

from facenet_pytorch import MTCNN, InceptionResnetV1

class EndpointHandler():
    def __init__(self, path=""):
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        self.mtcnn = MTCNN(device=self.device)
        self.resnet = InceptionResnetV1(pretrained='vggface2', device=self.device).eval()

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        imageData = data.pop("image", data)
        return imageData
        # image = Image.open(BytesIO(base64.b64decode(imageData)))
        # face_batch = self.mtcnn([image])
        # face_batch = [i for i in face_batch if i is not None]
        # if face_batch:
        #     aligned = torch.stack(face_batch)
        #     if self.device.type == "cuda":
        #         aligned = aligned.to(self.device)

        #     embeddings = self.resnet(aligned).detach().cpu()
        #     return embeddings.tolist()
        # else: return None