from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
import io
import base64
from ultralytics import YOLO
import torch
from PIL import Image, ImageOps

import utils
from drawing import draw_keypoints


device = 'cuda' if torch.cuda.is_available() else 'cpu'
print('Using device:', device)

model_pose = YOLO('yolov8l-pose.pt')
model_pose.to(device)

app = FastAPI()


@app.get("/")
async def health():
    return JSONResponse(content={"status": "ok"})


@app.post("/predict-image")
async def predict_image(file: UploadFile = File(...)):
    contents = await file.read()
    input_image = Image.open(io.BytesIO(contents)).convert("RGB")
    input_image = ImageOps.exif_transpose(input_image)

    # predict
    result = model_pose(input_image)[0]
    keypoints = utils.get_keypoints(result)

    # draw keypoints
    output_image = draw_keypoints(input_image, keypoints).convert("RGB")

    # calculate angles
    lea, rea = utils.get_eye_angles(keypoints)
    lba, rba = utils.get_elbow_angles(keypoints)
    angles  = {'left_eye_angle': lea, 'right_eye_angle': rea, 'left_elbow_angle': lba, 'right_elbow_angle': rba}

    # encode to base64
    img_buffer = io.BytesIO()
    output_image.save(img_buffer, format="JPEG")
    img_buffer.seek(0)
    img_base64 = base64.b64encode(img_buffer.getvalue()).decode("utf-8")

    # prepare json response
    json_data = {
        "keypoints": keypoints,
        "angles": angles,
        "output_image": img_base64
    }

    return JSONResponse(content=json_data)


# if __name__ == "__main__":
#     import uvicorn
#     uvicorn.run(app, host="0.0.0.0", port=7860)