import os
import re
import shutil

from config import MODEL_OUTPUT_IMAGE_NAMES

ROOT_DIR = "benchmark_images_generations"
OUTPUT_DIR = "data/"

def main():
    for domain in os.listdir(ROOT_DIR):
        domain_dir = os.path.join(ROOT_DIR, domain)
        for i, sample_dir in enumerate(os.listdir(domain_dir)):
            if sample_dir == ".DS_Store":
                continue
            sample_dir_path = os.path.join(domain_dir, sample_dir)
            prompt = sample_dir[4:].strip()

            output_sample_dir = os.path.join(OUTPUT_DIR, domain, f"sample_{i}")
            os.makedirs(output_sample_dir, exist_ok=True)

            # --- Look for all the images needed to complete the sample ---
            
            # Input background
            ## look for an image (either .jpg or .png) whose name is composed as "bg{number}"
            ## and copy it to the output directory, then rename it to input_bg.jpg
            input_bg = None
            for file in os.listdir(sample_dir_path):
                if re.match(r"bg\d+\.(jpg|png)", file):
                    input_bg = file
                    break
            if input_bg:
                input_bg_path = os.path.join(sample_dir_path, input_bg)
                shutil.copy(input_bg_path, os.path.join(output_sample_dir, "input_bg.jpg"))
            else:
                print(f"Warning: No input background found in {sample_dir_path}. Skipping sample {i}...")
                continue

            # Input foreground
            ## look for an image (either .jpg or .png) whose name is either composed as "fg_{alphanumeric}.jpg" or "fg_{alphanumeric}.png",
            ## or "fg{number}_{alphanumeric}.jpg" or "fg{number}_{alphanumeric}.png", and never contains the word "mask"
            ## and copy it to the output directory, then rename it to input_fg.jpg
            input_fg = None
            for file in os.listdir(sample_dir_path):
                if re.match(r"fg(_\w+)?\.(jpg|png)", file) or re.match(r"fg\d+(_\w+)?\.(jpg|png)", file):
                    if "mask" not in file:
                        input_fg = file
                        break
            if input_fg:
                input_fg_path = os.path.join(sample_dir_path, input_fg)
                shutil.copy(input_fg_path, os.path.join(output_sample_dir, "input_fg.jpg"))
            else:
                print(f"Warning: No input foreground found in {sample_dir_path}. Skipping sample {i}...")
                continue

            # Outputs
            if not all(
                [
                    os.path.exists(os.path.join(sample_dir_path, "cp_bg_fg.jpg")),
                    os.path.exists(os.path.join(sample_dir_path, "kvedit.jpg")),
                    os.path.exists(os.path.join(sample_dir_path, "tf-icon.png")),
                ]
            ) or not any(
                [
                    os.path.exists(os.path.join(sample_dir_path, "alphanoise0.05_timesteps50_QTrue_KTrue_VTrue_taua0.4_taub0.8_guidance3.0.png")),
                    os.path.exists(os.path.join(sample_dir_path, "alphanoise0.05_timesteps50_QTrue_KTrue_VTrue_taua0.4_taub0.8_guidance3.0_vital-layers.png")),
                ]
            ):
                print(f"Warning: Not all output images found in {sample_dir_path}. Skipping sample {i}...")
                # Delete the output directory for this sample
                shutil.rmtree(output_sample_dir)
                continue

            for model_name, image_name in MODEL_OUTPUT_IMAGE_NAMES.items():
                # Patch needed because of inconsistency in the naming of the dit-editor images
                if model_name != "dit-editor":
                    image_path = os.path.join(sample_dir_path, image_name)
                else:
                    if os.path.exists(os.path.join(sample_dir_path, "alphanoise0.05_timesteps50_QTrue_KTrue_VTrue_taua0.4_taub0.8_guidance3.0.png")):
                        image_path = os.path.join(sample_dir_path, "alphanoise0.05_timesteps50_QTrue_KTrue_VTrue_taua0.4_taub0.8_guidance3.0.png")
                    elif os.path.exists(os.path.join(sample_dir_path, "alphanoise0.05_timesteps50_QTrue_KTrue_VTrue_taua0.4_taub0.8_guidance3.0_vital-layers.png")):
                        image_path = os.path.join(sample_dir_path, "alphanoise0.05_timesteps50_QTrue_KTrue_VTrue_taua0.4_taub0.8_guidance3.0_vital-layers.png")
                    else:
                        raise FileNotFoundError(f"Neither of the expected images for dit-editor found in {sample_dir_path}")

                target_path = os.path.join(output_sample_dir, image_name)
                if os.path.exists(image_path):
                    shutil.copy(image_path, target_path)
                else:
                    print(f"Warning: {image_name} not found in {sample_dir_path}. Skipping...")

            # -- Create the prompt.txt file --
            ## create a file named prompt.txt in the output directory and write the prompt in it
            prompt_file_path = os.path.join(output_sample_dir, "prompt.txt")
            with open(prompt_file_path, "w") as prompt_file:
                prompt_file.write(prompt)

if __name__ == "__main__":
    main()