app.py

# app.py
import streamlit as st
import streamlit.components.v1 as components
import os
import json
import uuid
from PIL import Image, ImageDraw # For minimap
import urllib.parse # For decoding save data

# --- Constants ---
SAVE_DIR = "saved_worlds"
WORLD_MAP_FILE = os.path.join(SAVE_DIR, "world_map.json")
DEFAULT_SPACE_SIZE = 50 # Matches ground plane size in JS (approx)
MINIMAP_CELL_SIZE = 10 # Pixels per cell in minimap image

# --- Ensure Save Directory Exists ---
os.makedirs(SAVE_DIR, exist_ok=True)

# --- Helper Functions for Persistence ---

def load_world_map():
    """Loads the world map metadata."""
    if os.path.exists(WORLD_MAP_FILE):
        try:
            with open(WORLD_MAP_FILE, 'r') as f:
                return json.load(f)
        except json.JSONDecodeError:
            st.error("Error reading world map file. Starting fresh.")
            return {"spaces": {}} # Return empty if corrupt
    return {"spaces": {}} # SpaceID -> {"grid_x": int, "grid_y": int, "name": str}

def save_world_map(world_data):
    """Saves the world map metadata."""
    with open(WORLD_MAP_FILE, 'w') as f:
        json.dump(world_data, f, indent=4)

def save_space_data(space_id, objects_data):
    """Saves the object data for a specific space."""
    file_path = os.path.join(SAVE_DIR, f"{space_id}.json")
    with open(file_path, 'w') as f:
        # Store objects directly, could add metadata later
        json.dump({"objects": objects_data}, f, indent=4)

def load_space_data(space_id):
    """Loads object data for a specific space."""
    file_path = os.path.join(SAVE_DIR, f"{space_id}.json")
    if os.path.exists(file_path):
        try:
            with open(file_path, 'r') as f:
                data = json.load(f)
                return data.get("objects", []) # Return objects list or empty
        except json.JSONDecodeError:
            st.error(f"Error reading space file {space_id}.json.")
            return []
    return [] # Return empty list if file doesn't exist

def find_next_available_grid_slot(world_data):
    """Finds the next empty slot in a spiral pattern (simple version)."""
    occupied = set((d["grid_x"], d["grid_y"]) for d in world_data.get("spaces", {}).values())
    x, y = 0, 0
    dx, dy = 0, -1
    steps = 0
    limit = 1
    count = 0
    while (x, y) in occupied:
        if x == y or (x < 0 and x == -y) or (x > 0 and x == 1-y):
             dx, dy = -dy, dx # Change direction (spiral)
        x, y = x + dx, y + dy
        count += 1
        if count > 1000: # Safety break
             st.error("Could not find empty grid slot easily!")
             return None, None
    return x, y

# --- Minimap Generation ---
def generate_minimap(world_data, current_space_id=None):
    spaces = world_data.get("spaces", {})
    if not spaces:
        return None # No map if no spaces saved

    coords = [(d["grid_x"], d["grid_y"]) for d in spaces.values()]
    min_x = min(c[0] for c in coords)
    max_x = max(c[0] for c in coords)
    min_y = min(c[1] for c in coords)
    max_y = max(c[1] for c in coords)

    img_width = (max_x - min_x + 1) * MINIMAP_CELL_SIZE
    img_height = (max_y - min_y + 1) * MINIMAP_CELL_SIZE

    img = Image.new('RGB', (img_width, img_height), color = 'lightgrey')
    draw = ImageDraw.Draw(img)

    for space_id, data in spaces.items():
        cell_x = (data["grid_x"] - min_x) * MINIMAP_CELL_SIZE
        cell_y = (data["grid_y"] - min_y) * MINIMAP_CELL_SIZE
        color = "blue"
        if space_id == current_space_id:
            color = "red" # Highlight current space
        elif current_space_id is None and space_id == list(spaces.keys())[0]: # Highlight first if none selected
             color = "red"

        draw.rectangle(
            [cell_x, cell_y, cell_x + MINIMAP_CELL_SIZE -1, cell_y + MINIMAP_CELL_SIZE -1],
            fill=color, outline="black"
        )

    return img


# --- Page Config ---
st.set_page_config(
    page_title="Multiplayer World Builder",
    layout="wide"
)

# --- Initialize Session State ---
if 'selected_object' not in st.session_state:
    st.session_state.selected_object = 'None'
if 'current_space_id' not in st.session_state:
    st.session_state.current_space_id = None # Will be set when loading/creating
if 'space_name' not in st.session_state:
     st.session_state.space_name = ""
if 'initial_objects' not in st.session_state:
     st.session_state.initial_objects = [] # Objects to load into JS

# --- Load initial world data ---
world_data = load_world_map()

# --- Handle Save Data from JS (Query Param Workaround) ---
query_params = st.query_params.to_dict()
save_data_encoded = query_params.get("save_data")

save_triggered = False
if save_data_encoded:
    try:
        save_data_json = urllib.parse.unquote(save_data_encoded[0]) # Get first value if list
        objects_to_save = json.loads(save_data_json)

        space_id_to_save = query_params.get("space_id", [st.session_state.current_space_id])[0] # Get from URL or state
        space_name_to_save = query_params.get("space_name", [st.session_state.space_name])[0]

        if not space_id_to_save:
            space_id_to_save = str(uuid.uuid4()) # Create new ID
            st.session_state.current_space_id = space_id_to_save # Update state
            grid_x, grid_y = find_next_available_grid_slot(world_data)
            if grid_x is not None:
                 world_data.setdefault("spaces", {})[space_id_to_save] = {
                     "grid_x": grid_x,
                     "grid_y": grid_y,
                     "name": space_name_to_save or f"Space {len(world_data.get('spaces',{}))+1}"
                 }
                 save_world_map(world_data)
            else:
                 st.error("Failed to assign grid position!")


        # Save the actual object data
        save_space_data(space_id_to_save, objects_to_save)
        st.success(f"Space '{space_name_to_save or space_id_to_save}' saved successfully!")

        # Update name in world map if it changed and space exists
        if space_id_to_save in world_data.get("spaces", {}) and space_name_to_save:
             world_data["spaces"][space_id_to_save]["name"] = space_name_to_save
             save_world_map(world_data)

        # IMPORTANT: Clear query param to prevent resave on refresh
        st.query_params.clear()
        save_triggered = True # Flag to maybe skip immediate rerun if needed below

    except Exception as e:
        st.error(f"Error processing save data: {e}")
        st.exception(e) # Show traceback

# Need to reload objects if just saved or loading a new space
if 'current_space_id' in st.session_state and st.session_state.current_space_id:
     st.session_state.initial_objects = load_space_data(st.session_state.current_space_id)


# --- Sidebar Controls ---
with st.sidebar:
    st.title("🏗️ World Controls")

    # --- Space Management ---
    st.subheader("Manage Spaces")
    saved_spaces = list(world_data.get("spaces", {}).items()) # List of (id, data) tuples
    space_options = {sid: data.get("name", f"Unnamed ({sid[:6]}...)") for sid, data in saved_spaces}
    space_options["_new_"] = "✨ Create New Space ✨" # Special option

    selected_space_display = st.selectbox(
        "Load or Create Space:",
        options = ["_new_"] + list(space_options.keys()),
        format_func = lambda x: space_options.get(x, "Select...") if x != "_new_" else "✨ Create New Space ✨",
        index=0, # Default to Create New
        key="space_selection_key" # Unique key might be needed if dynamically changing options
        # Note: Changing this will trigger rerun. Need logic below to handle load.
    )

    # Handle Load/Create based on selection
    if st.session_state.space_selection_key != "_new_":
        # Load existing selected
        if st.session_state.current_space_id != st.session_state.space_selection_key:
             st.session_state.current_space_id = st.session_state.space_selection_key
             st.session_state.initial_objects = load_space_data(st.session_state.current_space_id)
             st.session_state.space_name = world_data["spaces"][st.session_state.current_space_id].get("name", "")
             st.rerun() # Rerun to load data and update JS injection
    elif st.session_state.space_selection_key == "_new_" and st.session_state.current_space_id is not None:
         # Handle switch from existing to "Create New"
         st.session_state.current_space_id = None
         st.session_state.initial_objects = []
         st.session_state.space_name = ""
         st.rerun()


    current_name = st.text_input(
         "Current Space Name:",
         value=st.session_state.space_name,
         key="current_space_name_input"
    )
    # Update state immediately if name changes
    if current_name != st.session_state.space_name:
        st.session_state.space_name = current_name
        # If editing an existing space, maybe update world_map.json immediately or on save? Let's do on save for simplicity.


    st.info(f"Current Space ID: {st.session_state.current_space_id or 'None (New)'}")
    st.caption("Saving uses the name above. A unique ID is assigned automatically for new spaces.")
    st.markdown("---")


    # --- Object Placement Controls ---
    st.subheader("Place Objects")
    object_types = ["None", "Simple House", "Tree", "Rock", "Fence Post"]
    selected = st.selectbox(
        "Select Object to Place:",
        options=object_types,
        key='selected_object'
    )
    st.markdown("---")

    # --- Minimap ---
    st.subheader("World Minimap")
    minimap_img = generate_minimap(world_data, st.session_state.current_space_id)
    if minimap_img:
        st.image(minimap_img, caption="Blue: Saved Spaces, Red: Current", use_column_width=True)
    else:
        st.caption("No spaces saved yet.")


# --- Main Area ---
st.header("3D Space Editor")
st.caption(f"Editing: {st.session_state.space_name or 'New Space'}")

# --- Load and Prepare HTML ---
html_file_path = 'index.html'

try:
    with open(html_file_path, 'r', encoding='utf-8') as f:
        html_template = f.read()

    # --- Inject Python state into JavaScript ---
    js_injection_script = f"""
<script>
    // Set global variables BEFORE the main script runs
    window.SELECTED_OBJECT_TYPE = {json.dumps(st.session_state.selected_object)};
    window.INITIAL_OBJECTS = {json.dumps(st.session_state.initial_objects)};
    window.CURRENT_SPACE_ID = {json.dumps(st.session_state.current_space_id)};
    window.CURRENT_SPACE_NAME = {json.dumps(st.session_state.space_name)};
    console.log("Streamlit State:", {{
        selectedObject: window.SELECTED_OBJECT_TYPE,
        initialObjects: window.INITIAL_OBJECTS,
        spaceId: window.CURRENT_SPACE_ID,
        spaceName: window.CURRENT_SPACE_NAME
    }});
</script>
"""
    # Insert the injection script just before the closing </head>
    # Using placeholder is safer if you modify index.html: html_content_with_state = html_template.replace('</head>', js_injection_script + '\n</head>', 1)


    # --- Embed HTML Component ---
    components.html(
        html_content_with_state,
        height=750, # Adjust height as needed
        scrolling=False
    )

except FileNotFoundError:
    st.error(f"Error: Could not find the file '{html_file_path}'.")
    st.warning(f"Please make sure `{html_file_path}` is in the same directory as `app.py` and that the `{SAVE_DIR}` directory exists.")
except Exception as e:
    st.error(f"An error occurred: {e}")
    st.exception(e)