import json
from pathlib import Path
from utils import generate, GEMINI_2DOT5_PRO_PREVIEW_05_06

from bedroom_prompts import (
    # BEDROOM SPECIFIC LAYOUT STYLE DESCRIPTIONS
    master_bedroom_traditional_spec,
    small_bedroom_efficient_spec,
    bedroom_with_workstation_spec,
    guest_bedroom_welcoming_spec,

    # EXAMPLE BEDROOM LAYOUTS
    rectangular_master_bedroom_traditional_example,
    l_shaped_bedroom_with_workstation_example,
    open_concept_small_bedroom_efficient_example,

    # ADAPTED PROMPT FOR BEDROOM LAYOUTS
    bedroom_layout_generation_prompt,

    # PROMPT FOR CHOOSING A BEDROOM LAYOUT STYLE
    bedroom_layout_style_choice_prompt
)

def pick_right_spec(layout_style):
    if 'master' in layout_style.lower():
        return master_bedroom_traditional_spec
    elif 'small' in layout_style.lower():
        return small_bedroom_efficient_spec
    elif 'workstation' in layout_style.lower():
        return bedroom_with_workstation_spec
    elif 'guest' in layout_style.lower():
        return guest_bedroom_welcoming_spec
    else:
        print(f"WARNING: Unknown layout style: {layout_style}")
        return ""

def pick_one_shot_example(shape):
    if 'open' in shape.lower():
        return open_concept_small_bedroom_efficient_example
    if 'l-shape' in shape.lower():
        return l_shaped_bedroom_with_workstation_example
    elif 'rectangular' in shape.lower():
        return rectangular_master_bedroom_traditional_example
    else:
        print(f"WARNING: Unknown room shape: {shape}")
        return ""
    
# python run_bedroom_layout_generation.py
def main():
    # Logs
    bedroom_styles = []
    layout_prompts = []
    generated_layouts = []

    model_id = GEMINI_2DOT5_PRO_PREVIEW_05_06

    for i in range(5, 55):

        try:
            input_json_file_path = "bedrooms_final/generated_bedrooms_{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_bedroom_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
                        bedroom_style_answer = generate(bedroom_layout_style_choice_prompt.format(room_input=room), model_id)
                        chosen_bedroom_style = json.loads(bedroom_style_answer.strip().removeprefix('```json').removesuffix('```'))
                        bedroom_styles.append(chosen_bedroom_style)
                        print(chosen_bedroom_style['chosen_layout_style'])
                        
                        # Layout generation
                        layout_prompt = bedroom_layout_generation_prompt.format(
                            room_details=room,
                            layout_style=chosen_bedroom_style,
                            bedroom_layout_specifics=pick_right_spec(chosen_bedroom_style['chosen_layout_style']),
                            bedroom_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()