import json
from living_room_prompts import living_room_generation_prompt
from utils import draw_room, generate, GEMINI_2DOT5_PRO_PREVIEW_05_06

from pathlib import Path

def main():
    model_id = GEMINI_2DOT5_PRO_PREVIEW_05_06

    images_output_dir = "living_rooms/images"
    base_file_name = "generated_living_rooms"

    output_dir = Path("living_rooms_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 = 10
    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 living rooms, iteration {i}, variant {v}...")

                try:
                    rooms = generate(living_room_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")
                    print(e)
            else:
                print(f"Living rooms, iteration {i}, variant {v} already exist. Skip.")

if __name__ == "__main__":
    main()