import os
from PIL import Image

class FileProcessor:
    PARTS_NAMES = [
        "Quarter-panel", "Front-wheel", "Back-window", "Trunk", "Front-door", "Rocker-panel",
        "Grille", "Windshield", "Front-window", "Back-door", "Headlight", "Back-wheel",
        "Back-windshield", "Hood", "Fender", "Tail-light", "License-plate", "Front-bumper",
        "Back-bumper", "Mirror", "Roof"
    ]

    DAMAGE_NAMES = [
        "Broken part", "Scratch", "Dent", "Paint chip"
    ]

    @staticmethod
    def convert_normalized_to_original(coords, width, height):
        original_coords = []
        for i in range(0, len(coords), 2):
            x = coords[i] * width
            y = coords[i + 1] * height
            original_coords.extend([x, y])
        return original_coords

    def process_file(self, file_path, class_names, image_size):
        width, height = image_size
        with open(file_path, "r") as f:
            lines = f.readlines()
        
        new_lines = []
        for line in lines:
            parts = line.strip().split()
            class_index = int(parts[0])
            coords = list(map(float, parts[1:]))
            original_coords = self.convert_normalized_to_original(coords, width, height)
            class_name = class_names[class_index]
            new_line = f"{class_name} " + " ".join(map(str, original_coords)) + "\n"
            new_lines.append(new_line)
        
        with open(file_path, "w") as f:
            f.writelines(new_lines)

    def process_folder(self, folder_path, class_names):
        labels_folder = os.path.join(folder_path, "labels")
        for file in os.listdir(folder_path):
            if file.endswith(".jpg"):
                image_path = os.path.join(folder_path, file)
                if os.path.exists(image_path):
                    image = Image.open(image_path)
                    image_size = image.size
                    txt_file = file.replace(".jpg", ".txt")
                    txt_file_path = os.path.join(labels_folder, txt_file)
                    if os.path.exists(txt_file_path):
                        self.process_file(txt_file_path, class_names, image_size)
                    else:
                        print(f"Text file {txt_file_path} not found.")
                else:
                    print(f"Image {image_path} not found.")

    def process_output_folder(self, output_folder):
        parts_folder = os.path.join(output_folder, "parts")
        damage_folder = os.path.join(output_folder, "damage")

        self.process_folder(parts_folder, self.PARTS_NAMES)
        self.process_folder(damage_folder, self.DAMAGE_NAMES)