import glob
import os
import subprocess
from loguru import logger
import shutil
from fastapi import FastAPI
from schemas.api_schema import optimization
import uvicorn

app = FastAPI()

class ModelOptimization:
    def __init__(self, num_class, mount_path, weight_path, image_size=416):
        self.num_class = num_class
        self.image_size = image_size
        self.mount_path = mount_path
        self.weight_path = weight_path

    def change_configurations(self):
        logger.info(f"Provided number of classes and image size are :  {self.num_class} and {self.image_size}")
        try:

            with open('tensorrtx/yolov5/yololayer.h', 'r') as file:
                # read a list of lines into data
                data = file.readlines()

            data[19] = f"    static constexpr int CLASS_NUM = {self.num_class};\n"
            data[20] = f"    static constexpr int INPUT_H = {self.image_size};\n"
            data[21] = f"    static constexpr int INPUT_W = {self.image_size};\n"

            # and write everything back
            with open('tensorrtx/yolov5/yololayer.h', 'w') as file:
                file.writelines(data)
            logger.info("Successfully changed configurations")
        except Exception as e:
            logger.info(f"Failed to change configurations :  {e}")
    def optimize_model(self):
        try:
            code_directory = "/opt/yolo_model_optimization"
            logger.info(f"Current directory is :  {code_directory}")
            if os.path.exists(os.path.join(code_directory, 'tensorrtx/yolov5/build')):
                shutil.rmtree(os.path.join(code_directory, 'tensorrtx/yolov5/build'))
                logger.info("Removed existing build directory")
            os.mkdir(os.path.join(code_directory, 'tensorrtx/yolov5/build'))
            logger.info("Built new build directory")
            shutil.copy(self.weight_path, os.path.join(code_directory, 'tensorrtx/yolov5/build'))
            logger.info("Copied wts file to build directory")
            weight_name_with_extension = os.path.basename(self.weight_path)
            os.chdir(os.path.join(code_directory, 'tensorrtx/yolov5/build'))
            logger.info("Changed to build directory")
            logger.info("Running CMake command")
            subprocess.run(['cmake', '..'])
            logger.info("Running Make command")
            subprocess.run(['make'])
            logger.info("Optimizing model")
            engine_name = "best.engine"
            subprocess.run(["sudo", "./yolov5", "-s",os.path.join(code_directory, 'tensorrtx/yolov5/build', weight_name_with_extension) , engine_name, "c", "0.33", "0.50"])
            shutil.move(os.path.join(code_directory, 'tensorrtx/yolov5/build', engine_name), self.mount_path)
            shutil.move(os.path.join(code_directory, 'tensorrtx/yolov5/build', "libmyplugins.so"), self.mount_path)
            logger.info(f"Moved engine file and libmyplugins.so file to {self.mount_path}")

        except Exception as e:
            logger.info(f"Failed to optimized model :  {e}")

@app.post("/optimize")
async def root(content: optimization):
    obj = ModelOptimization(num_class=int(content.num_class), image_size=int(content.image_size), weight_path = content.weight_path, mount_path=content.mount_path)
    obj.change_configurations()
    obj.optimize_model()
    return {"message": "successfull"}


if __name__ == '__main__':
    uvicorn.run(app, port=8080, host='0.0.0.0')


