from fastapi import FastAPI, File, UploadFile
import deepface.DeepFace as DeepFace
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import numpy as np
from mivolo.predictor import Predictor
import base64
from PIL import Image
import numpy as np
from AttributesHolder import Namespace
import cv2

config = {
    "checkpoint": 'models/mivolo_imbd.pth.tar',
    "detector_weights": 'models/yolov8x_person_face.pt',
    "device": 'cpu',
    "draw": False,
    "with_persons": True,
    "disable_faces": False,
    "output": 'output'
}

namespace = Namespace()
setattr(namespace, 'checkpoint', 'models/mivolo_imbd.pth.tar')
setattr(namespace, 'detector_weights', 'models/yolov8x_person_face.pt')
setattr(namespace, 'device', 'cpu')
setattr(namespace, 'draw', False)
setattr(namespace, 'with_persons', True)
setattr(namespace, 'disable_faces', False)
setattr(namespace, 'output', 'output')

predictor = Predictor(config=namespace)

app = FastAPI()

class Base64Data(BaseModel):
    base64_data: str


origins = ["*"]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
@app.get("/")
def index():
    return {"message": "Hello, World!!!"}

# post image to server
@app.post("/predict-actor/")
async def create_upload_file(contents: Base64Data):
    try:
        # Read the file
        print("-----Starting Prediction----------")
        loaded_image = base64_to_cv2(contents.base64_data)
        detected_objects, out_im = predictor.recognize(loaded_image)
        age = detected_objects.ages[0]
        gender = detected_objects.genders[0]
        contents = contents.base64_data
        df = DeepFace.find(img_path = contents, db_path = "dataset/",model_name ='GhostFaceNet', threshold=0.9)
        
        filename = df[0].head()['identity'][0]
        filename = filename.replace("\\", "/")

        print(f"filename: {filename}")

        # Convert the image to base64
        base64_output = image_to_base64("dataset/" + filename)

        print("------Prediction Done-----------")
        return JSONResponse(content={
            "celeb_image": base64_output,
            "celeb":df[0].head()['identity'][0], "res":{
            "age": age,
            "gender": gender
        } }, status_code=200)
    except Exception as e:
        return JSONResponse(content={"message": "Error processing the file.", "error": str(e)}, status_code=500)    



def image_to_base64(image_path):
    with open(image_path, "rb") as img_file:
        # Read the image file
        img_data = img_file.read()
        # Encode the image data as base64
        base64_data = base64.b64encode(img_data)
        # Decode bytes-like object to ASCII string
        base64_str = base64_data.decode("ascii")
        return base64_str


def base64_to_cv2(base64_string):
    base64_string = base64_string.split(",")[1]
    # Decode the base64 string into bytes
    decoded_bytes = base64.b64decode(base64_string)
    
    # Convert bytes to numpy array
    np_array = np.frombuffer(decoded_bytes, np.uint8)
    
    # Decode the numpy array into an image
    image = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
    
    return image