import bentoml
from face_matcher import FaceMatcher
from pydantic import BaseModel, Field
from typing import List
import base64
import json

class FaceMatchResponse(BaseModel):
    confidence: float = Field(description="The confidence that the document and selfie pictures match.")
    raw: str = Field(description="Metadata for ML retraining and debugging.")

@bentoml.service(
    resources={"cpu": "1"},
    traffic={"timeout": 60},
    # ssl={
    #     "enabled": True,
    #     "certfile": "/home/aidan/janus/include/cert.pem",
    #     "keyfile": "/home/aidan/janus/include/key.pem",
    #     "keyfile_password": "vulcan",
    #     "version": 3,
    #     "ciphers": "TLSv1.3",
    # }
)
class FaceMatchService:
    # Wrapper around 
    def __init__(self) -> None:
        self.face_matcher = FaceMatcher()

    @bentoml.api
    def match_pair(self, image_document: str, image_selfie: str, return_vectors=True) -> FaceMatchResponse:
        """
        Inputs:
            image_document: ID image encoded as a base64 ascii string.
            image_selfie: Selfie image encoded as a base64 ascii string.
            return_vectors: Whether to return embedding vectors in response (default True)
        
        Returns: 
            confidence: Processed similarity score ranging from 0 to 1
            raw:
                similarity: Raw similarity score from -1 to 1
                id: Id metadata, including autocapture and blazeface outputs and embedding vector if return_vectors=True
                selfie: Selfie metadata, including blazeface outputs and embedding vector if return_vectors=True
        """
        id = base64.b64decode(image_document)
        selfie = base64.b64decode(image_selfie)
        payload = self.face_matcher.match_pair(id=id, selfie=selfie, return_vectors=return_vectors)
        response = FaceMatchResponse(confidence=payload["final_score"], raw=json.dumps(payload))
        return response

    # @bentoml.api(batchable=True, max_batch_size=16, max_latency_ms=1000)
    # def match(self, ids: t.List[Image.Image], selfies: t.List[Image.Image]):
    #     ids = [Image.open(io.BytesIO(id)).convert("RGB") for id in ids]
    #     selfies = [Image.open(io.BytesIO(selfie)).convert("RGB") for selfie in selfies]
    #     return self.face_matcher.match(ids, selfies)


# FaceEmbedder.add_asgi_middleware(HTTPSRedirectMiddleware)
