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

app = FastAPI()

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

    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_H = {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, weight_path):
        try:
            shutil.copy(weight_path, 'tensorrtx/yolov5/build')
            weight_name_with_extension = os.path.basename(weight_path)
            weight_name, extension = os.path.splitext(weight_name_with_extension)
            current_directory = os.getcwd()
            logger.info(f"Current directory is :  {current_directory}")
            build_files = glob.glob("tensorrtx/yolov5/build/*")
            for file in build_files:
                os.remove(file)

            # build_path = os.path.join(current_directory, "yolov5", "build")
            # os.mkdir(build_path)
            # logger.info(f"Created build folder")
            os.chdir('tensorrtx/yolov5/build')
            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", weight_name_with_extension, engine_name, "c", "0.33", "0.50"])

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

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


