from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import asyncio
import uvicorn

import uuid
import time 
import sdtcloudpubsub
import logging.handlers
import glob
import traceback
# Communication use grpc
import grpc
import grpc.aio
import utils.image_pb2 as pb2
import utils.image_pb2_grpc as pb2_grpc

###############################################
#               Logger Setting                #
###############################################
logger = logging.getLogger()
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')

log_fileHandler = logging.handlers.RotatingFileHandler(
    filename=f"./logs/log_request.log",
    maxBytes=1024000,
    backupCount=3,
    mode='a')

log_fileHandler.setFormatter(formatter)
logger.addHandler(log_fileHandler)

###############################################
#                MQTT Setting                 #
###############################################
sdtcloud = sdtcloudpubsub.sdtcloudpubsub()
sdtcloud.setClient(f"device-app-{uuid.uuid1()}") # parameter is client ID(string)

###############################################
#              FastAPI instance               #
###############################################
app = FastAPI()

###############################################
#                 grpc client                 #
###############################################

async def upload_image(stub, filename):
    with open(filename, 'rb') as f:
        image_data = f.read()
    response = await stub.UploadImage(pb2.ImageRequest(filename=filename, image_data=image_data))
    print(response.message)
    print(f"Inference result: {response.inference_result}")
    return response

async def run(image_path):
    async with grpc.aio.insecure_channel('localhost:50051') as channel:
        stub = pb2_grpc.ImageServiceStub(channel)
        response = await upload_image(stub, image_path)
        return response


class ImagePathRequest(BaseModel):
    image_path: str = "./data/test_1.jpg"

# FastAPI endpoint
@app.post("/upload_image/")
async def upload_image_endpoint(request: ImagePathRequest):
    start_time = time.time()
    try:
        response = await run(request.image_path)
        logger.info('Successfully requested!')
        end_time = time.time()
        data = {
            "processingTime": end_time - start_time
        }

        print(f"#### Topic: {sdtcloud.topic}")
        print(f"#### DATA: {data}")
        sdtcloud.pubReqMessage(data)

        return {
            "message": response.message,
            "inference_result": response.inference_result
        }
    except Exception as e:
        logger.error(traceback.format_exc())
        raise HTTPException(status_code=500, detail="Internal Server Error")

if __name__ == '__main__':    
    uvicorn.run("main:app", host="0.0.0.0", port=8888, reload=True)