import json
from bedroom_prompts import bedroom_generation_prompt  # add prompts to the path
from utils import generate, GEMINI_2DOT5_PRO_PREVIEW_05_06  # add .. to the path

from pathlib import Path


def main():
    model_id = GEMINI_2DOT5_PRO_PREVIEW_05_06

    base_file_name = "generated_bedrooms"

    output_dir = Path("bedrooms_final")  #  / "{:05d}".format(i) / "{:05d}".format(v))
    output_dir.mkdir(parents=True, exist_ok=True)

    num_rooms = 1000
    num_runs = num_rooms // 20
    num_variants = 1

    start_idx = 5
    for i in range(start_idx, start_idx + num_runs):

        for v in range(num_variants):
            output_file_path = Path(f"{output_dir}/{base_file_name}_{i}_variant_{v}.json")

            if not output_file_path.is_file():
                with open(output_file_path, "w") as file:
                    pass

                print(f"Generating bedrooms, iteration {i}, variant {v}...")

                try:
                    # Generate the bedrooms
                    rooms = generate(bedroom_generation_prompt, model_id)

                    processed_rooms = json.loads(rooms.strip().removeprefix("```json").removesuffix("```"))

                    # Save the processed_rooms to a JSON file
                    with open(output_file_path, "w") as file:
                        json.dump(processed_rooms, file, indent=4)
                except Exception as e:
                    print("ERROR", e)
            else:
                print(f"Bedrooms, iteration {i}, variant {v} already exist. Skip.")


if __name__ == "__main__":
    main()
