from google import genai
from google.genai.types import Tool, GenerateContentConfig, GoogleSearch
from google.genai import types

import json
import os
import time

import importlib

import utils
importlib.reload(utils)

from pathlib import Path
from utils import generate, GEMINI_2DOT5_PRO_PREVIEW_05_06 # add .. to the path

import living_room_prompts
importlib.reload(living_room_prompts)

from living_room_prompts import (
    # LIVING ROOM SPECIFIC LAYOUT STYLE DESCRIPTIONS
    conversational_living_room_spec,
    fireplace_centric_living_room_spec,
    multi_zone_living_room_spec,
    tv_focused_living_room_spec,

    # EXAMPLE LIVING ROOM LAYOUTS
    l_shaped_living_room_conversational_example,
    rectangular_living_room_tv_focused_example,
    open_living_room_tv_focused_example,

    # ADAPTED PROMPT FOR LIVING ROOM LAYOUTS
    living_room_layout_generation_prompt,

    # PROMPT FOR CHOOSING A LIVING ROOM LAYOUT STYLE
    living_room_layout_style_choice_prompt
) # add prompts to the path

def pick_right_spec(layout_style):
    if 'conversational' in layout_style.lower():
        return conversational_living_room_spec
    elif 'fireplace-centric' in layout_style.lower():
        return fireplace_centric_living_room_spec
    elif 'multi-zone' in layout_style.lower():
        return multi_zone_living_room_spec
    elif 'tv-focused' in layout_style.lower():
        return tv_focused_living_room_spec
    else:
        print(f"WARNING: Unknown layout style: {layout_style}")
        return ""

def pick_one_shot_example(shape):
    if 'open' in shape.lower():
        return open_living_room_tv_focused_example
    if 'l-shape' in shape.lower():
        return l_shaped_living_room_conversational_example
    elif 'rectangular' in shape.lower():
        return rectangular_living_room_tv_focused_example
    else:
        print(f"WARNING: Unknown room shape: {shape}")
        return ""
    

def main():
    # Logs
    living_room_styles = []
    layout_prompts = []
    generated_layouts = []

    model_id = GEMINI_2DOT5_PRO_PREVIEW_05_06

    for i in range(10, 60):

        try:
            input_json_file_path = "living_rooms_final/generated_living_rooms_{0}_variant_0.json".format(i)
            with open(input_json_file_path, "r") as file:
                processed_rooms = json.load(file)

            print(model_id)

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

            num_variants = 1 # 3

            for room in processed_rooms: # [10:]:
                room['layout_id'] = room['layout_id'] # + i * 20

                for v in range(num_variants):
                    print(f"Processed room {room['layout_id']}")

                    output_file_path = output_dir / f"upd_one_shot_room_{room['layout_id']}_variant_{v}.json"

                    if output_file_path.is_file():
                        continue

                    with open(output_file_path, "w") as file:
                        pass

                    try:
                        # Counter type choice
                        living_room_style_answer = generate(living_room_layout_style_choice_prompt.format(room_input=room), model_id)
                        chosen_living_room_style = json.loads(living_room_style_answer.strip().removeprefix('```json').removesuffix('```'))
                        living_room_styles.append(chosen_living_room_style)
                        print(chosen_living_room_style['chosen_layout_style'])
                        
                        # Layout generation
                        layout_prompt = living_room_layout_generation_prompt.format(
                            room_details=room,
                            layout_style=chosen_living_room_style,
                            living_room_layout_specifics=pick_right_spec(chosen_living_room_style['chosen_layout_style']),
                            living_room_layout_example=pick_one_shot_example(room['room']['shape'])
                        )
                        layout_prompts.append(layout_prompt)

                        layout_answer = generate(layout_prompt, model_id)
                        generated_room = json.loads(layout_answer.strip().removeprefix('```json').removesuffix('```'))
                        generated_layouts.append(generated_room)

                        with open(output_file_path, "w") as file:
                            json.dump(generated_room, file, indent=4)

                        print("Ok. Next.")
                    except Exception as e:
                        print(f"Error processing room {room['layout_id']}: {e}")
        except Exception as e:
            print(f"Error loading JSON File {input_json_file_path}: {e}")
    

if __name__ == "__main__":
    main()