app.py

# app.py
import streamlit as st

# --- Must be the first Streamlit command ---
st.set_page_config(page_title="Infinite World Builder", layout="wide")

import streamlit.components.v1 as components
import os
import json
import pandas as pd
import uuid
import math
import time

# Import our GameState class from gamestate.py
from gamestate import GameState

# --- Constants ---
SAVE_DIR = "saved_worlds"
GLOBAL_SAVES_DIR = "global_saves"  # Folder for global game saves
PLOT_WIDTH = 50.0  # Width of each plot in 3D space
PLOT_DEPTH = 50.0  # Depth of each plot
CSV_COLUMNS = ['obj_id', 'type', 'pos_x', 'pos_y', 'pos_z', 'rot_x', 'rot_y', 'rot_z', 'rot_order']

# --- Ensure Required Directories Exist ---
os.makedirs(SAVE_DIR, exist_ok=True)
os.makedirs(GLOBAL_SAVES_DIR, exist_ok=True)

@st.cache_data(ttl=3600)
def load_plot_metadata():
    """Scans SAVE_DIR for plot files and returns metadata."""
    plots = []
    try:
        plot_files = [f for f in os.listdir(SAVE_DIR) if f.endswith(".csv") and f.startswith("plot_X")]
    except FileNotFoundError:
        st.error(f"Save directory '{SAVE_DIR}' not found.")
        return []
    except Exception as e:
        st.error(f"Error listing save directory '{SAVE_DIR}': {e}")
        return []

    parsed_plots = []
    for filename in plot_files:
        try:
            parts = filename[:-4].split('_')  # Remove .csv extension
            grid_x = int(parts[1][1:])  # After 'X'
            grid_z = int(parts[2][1:])  # After 'Z'
            plot_name = " ".join(parts[3:]) if len(parts) > 3 else f"Plot ({grid_x},{grid_z})"
            parsed_plots.append({
                'id': filename[:-4],
                'filename': filename,
                'grid_x': grid_x,
                'grid_z': grid_z,
                'name': plot_name,
                'x_offset': grid_x * PLOT_WIDTH,
                'z_offset': grid_z * PLOT_DEPTH
            })
        except (IndexError, ValueError):
            st.warning(f"Could not parse grid coordinates from filename: {filename}. Skipping.")
            continue

    parsed_plots.sort(key=lambda p: (p['grid_x'], p['grid_z']))
    return parsed_plots

def load_plot_objects(filename, x_offset, z_offset):
    """Loads objects from a CSV file and applies world offsets."""
    file_path = os.path.join(SAVE_DIR, filename)
    objects = []
    try:
        df = pd.read_csv(file_path)
        if not all(col in df.columns for col in ['type', 'pos_x', 'pos_y', 'pos_z']):
            st.warning(f"CSV '{filename}' missing essential columns. Skipping.")
            return []
        df['obj_id'] = df.get('obj_id', pd.Series([str(uuid.uuid4()) for _ in range(len(df))]))
        for col, default in [('rot_x', 0.0), ('rot_y', 0.0), ('rot_z', 0.0), ('rot_order', 'XYZ')]:
            if col not in df.columns:
                df[col] = default

        for _, row in df.iterrows():
            obj_data = row.to_dict()
            obj_data['pos_x'] += x_offset
            obj_data['pos_z'] += z_offset
            objects.append(obj_data)
        return objects
    except FileNotFoundError:
        st.error(f"File not found during object load: {filename}")
        return []
    except pd.errors.EmptyDataError:
        return []
    except Exception as e:
        st.error(f"Error loading objects from {filename}: {e}")
        return []

def save_plot_data(filename, objects_data_list, plot_x_offset, plot_z_offset):
    """Saves object data list to a CSV file, making positions relative to the plot origin."""
    file_path = os.path.join(SAVE_DIR, filename)
    relative_objects = []
    if not isinstance(objects_data_list, list):
        st.error("Invalid data format received for saving (expected a list).")
        return False

    for obj in objects_data_list:
        pos = obj.get('position', {})
        rot = obj.get('rotation', {})
        obj_type = obj.get('type', 'Unknown')
        obj_id = obj.get('obj_id', str(uuid.uuid4()))

        if not all(k in pos for k in ['x', 'y', 'z']) or obj_type == 'Unknown':
            print(f"Skipping malformed object during save prep: {obj}")
            continue

        relative_obj = {
            'obj_id': obj_id, 
            'type': obj_type,
            'pos_x': pos.get('x', 0.0) - plot_x_offset,
            'pos_y': pos.get('y', 0.0),
            'pos_z': pos.get('z', 0.0) - plot_z_offset,
            'rot_x': rot.get('_x', 0.0), 
            'rot_y': rot.get('_y', 0.0),
            'rot_z': rot.get('_z', 0.0), 
            'rot_order': rot.get('_order', 'XYZ')
        }
        relative_objects.append(relative_obj)

    try:
        df = pd.DataFrame(relative_objects, columns=CSV_COLUMNS)
        df.to_csv(file_path, index=False)
        st.success(f"Saved {len(relative_objects)} objects to {filename}")
        return True
    except Exception as e:
        st.error(f"Failed to save plot data to {filename}: {e}")
        return False

# --- Initialize GameState Singleton ---
@st.cache_resource
def get_game_state():
    # This instance is shared across all sessions and reruns.
    return GameState(save_dir=SAVE_DIR, csv_filename="world_state.csv")

game_state = get_game_state()

# --- Session State Initialization ---
if 'selected_object' not in st.session_state:
    st.session_state.selected_object = 'None'
if 'new_plot_name' not in st.session_state:
    st.session_state.new_plot_name = ""
if 'js_save_data_result' not in st.session_state:
    st.session_state.js_save_data_result = None
if 'player_position' not in st.session_state:
    # This can be updated from JS if needed
    st.session_state.player_position = {"x": 0, "y": 0, "z": 0}
if 'loaded_global_state' not in st.session_state:
    st.session_state.loaded_global_state = None

plots_metadata = load_plot_metadata()
all_initial_objects = []
for plot in plots_metadata:
    all_initial_objects.extend(load_plot_objects(plot['filename'], plot['x_offset'], plot['z_offset']))

# --- Update GameState with initial objects if empty ---
if not game_state.get_state():
    game_state.update_state(all_initial_objects)

# --- Sidebar ---
with st.sidebar:
    st.title("🏗️ World Controls")
    st.header("Navigation (Plots)")
    st.caption("Click to teleport player to a plot.")
    max_cols = 2
    cols = st.columns(max_cols)
    col_idx = 0
    sorted_plots_for_nav = sorted(plots_metadata, key=lambda p: (p['grid_x'], p['grid_z']))
    for plot in sorted_plots_for_nav:
        button_label = f"➡️ {plot.get('name', plot['id'])} ({plot['grid_x']},{plot['grid_z']})"
        if cols[col_idx].button(button_label, key=f"nav_{plot['id']}"):
            target_x = plot['x_offset']
            target_z = plot['z_offset']
            try:
                from streamlit_js_eval import streamlit_js_eval
                js_code = f"teleportPlayer({target_x + PLOT_WIDTH/2}, {target_z + PLOT_DEPTH/2});"
                streamlit_js_eval(js_code=js_code, key=f"teleport_{plot['id']}")
            except Exception as e:
                st.error(f"Failed to send teleport command: {e}")
        col_idx = (col_idx + 1) % max_cols

    st.markdown("---")
    st.header("Place Objects")
    object_types = ["None", "Simple House", "Tree", "Rock", "Fence Post"]
    current_object_index = object_types.index(st.session_state.selected_object) if st.session_state.selected_object in object_types else 0
    selected_object_type_widget = st.selectbox("Select Object:", options=object_types, index=current_object_index, key="selected_object_widget")
    if selected_object_type_widget != st.session_state.selected_object:
         st.session_state.selected_object = selected_object_type_widget

    st.markdown("---")
    st.header("Save Work (Per Plot)")
    st.caption("Saves newly placed objects to the current plot. A new plot file is created for new areas.")
    if st.button("💾 Save Current Work", key="save_button"):
        from streamlit_js_eval import streamlit_js_eval
        js_get_data_code = "getSaveDataAndPosition();"
        streamlit_js_eval(js_code=js_get_data_code, key="js_save_processor")
        st.rerun()

    st.markdown("---")
    st.header("Global Save & Load")
    # Global Save: Persists the entire game state plus current player position.
    if st.button("💾 Global Save"):
        global_save_data = {
             "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
             "game_state": game_state.get_state(),
             "player_position": st.session_state.get("player_position", {"x": 0, "y": 0, "z": 0})
        }
        default_save_name = f"save_{time.strftime('%Y%m%d_%H%M%S')}.json"
        save_file_path = os.path.join(GLOBAL_SAVES_DIR, default_save_name)
        with open(save_file_path, "w", encoding="utf-8") as f:
             json.dump(global_save_data, f, indent=2)
        st.success(f"Global state saved to {default_save_name}")
        # Also store it in session state for download purposes.
        st.session_state.loaded_global_state = global_save_data

    # List global save files available in the folder.
    st.subheader("📂 Global Saves")
    save_files = sorted([f for f in os.listdir(GLOBAL_SAVES_DIR) if f.endswith(".json")])
    for file in save_files:
        if st.button(f"Load {file}", key=file):
            file_path = os.path.join(GLOBAL_SAVES_DIR, file)
            with open(file_path, "r", encoding="utf-8") as f:
                loaded_save = json.load(f)
            # Update the shared game state with loaded data.
            with game_state.lock:
                game_state.world_state = loaded_save.get("game_state", [])
            st.session_state.player_position = loaded_save.get("player_position", {"x": 0, "y": 0, "z": 0})
            st.session_state.loaded_global_state = loaded_save
            st.success(f"Global state loaded from {file}")

    st.markdown("---")
    st.header("Download Global Save as Markdown")
    # Use the loaded global state if available; otherwise use a default dictionary.
    current_save = st.session_state.get("loaded_global_state")
    if current_save is None:
        current_save = {"timestamp": "N/A", "game_state": [], "player_position": {"x": 0, "y": 0, "z": 0}}
    default_md_name = current_save.get("timestamp", "save").replace(":", "").replace(" ", "_") + ".md"
    download_name = st.text_input("Override File Name", value=default_md_name)
    if st.button("Generate Markdown & Download"):
        md_outline = f"""# Global Save: {download_name}
- ⏰ **Timestamp:** {current_save.get("timestamp", "N/A")}
- 🎮 **Number of Game Objects:** {len(current_save.get("game_state", []))}
- 🧭 **Player Position:** {current_save.get("player_position", {"x": 0, "y": 0, "z": 0})}

## Game Objects:
"""
        for i, obj in enumerate(current_save.get("game_state", []), start=1):
            obj_type = obj.get("type", "Unknown")
            pos = (obj.get("pos_x", 0), obj.get("pos_y", 0), obj.get("pos_z", 0))
            md_outline += f"- {i}. ✨ **{obj_type}** at {pos}\n"
        st.download_button("Download Markdown Save", data=md_outline, file_name=download_name, mime="text/markdown")

# --- Process Save Data from JS (Per Plot Save) ---
save_data_from_js = st.session_state.get("js_save_processor", None)
if save_data_from_js is not None:
    st.info("Received save data from client...")
    save_processed_successfully = False
    try:
        payload = json.loads(save_data_from_js) if isinstance(save_data_from_js, str) else save_data_from_js
        if isinstance(payload, dict) and 'playerPosition' in payload and 'objectsToSave' in payload:
            player_pos = payload['playerPosition']
            objects_to_save = payload['objectsToSave']
            if isinstance(objects_to_save, list):
                target_grid_x = math.floor(player_pos.get('x', 0.0) / PLOT_WIDTH)
                target_grid_z = math.floor(player_pos.get('z', 0.0) / PLOT_DEPTH)
                target_filename = f"plot_X{target_grid_x}_Z{target_grid_z}.csv"
                target_plot_x_offset = target_grid_x * PLOT_WIDTH
                target_plot_z_offset = target_grid_z * PLOT_DEPTH
                st.write(f"Attempting to save plot: {target_filename} (Player at: x={player_pos.get('x', 0):.1f}, z={player_pos.get('z', 0):.1f})")
                is_new_plot_file = not os.path.exists(os.path.join(SAVE_DIR, target_filename))
                save_ok = save_plot_data(target_filename, objects_to_save, target_plot_x_offset, target_plot_z_offset)
                if save_ok:
                    load_plot_metadata.clear()  # Clear cache so metadata reloads
                    try:
                        from streamlit_js_eval import streamlit_js_eval
                        streamlit_js_eval(js_code="resetNewlyPlacedObjects();", key="reset_js_state")
                    except Exception as js_e:
                        st.warning(f"Could not reset JS state after save: {js_e}")
                    if is_new_plot_file:
                        st.success(f"New plot created and saved: {target_filename}")
                    else:
                        st.success(f"Updated existing plot: {target_filename}")
                    # Update shared game state with new objects from this session
                    game_state.update_state(objects_to_save)
                    save_processed_successfully = True
                else:
                    st.error(f"Failed to save plot data to file: {target_filename}")
            else:
                st.error("Invalid 'objectsToSave' format received (expected list).")
        else:
            st.error("Invalid save payload structure received from client.")
    except json.JSONDecodeError:
        st.error("Failed to decode save data from client.")
    except Exception as e:
        st.error(f"Error processing save: {e}")
    st.session_state.js_save_processor = None
    if save_processed_successfully:
        st.rerun()

# --- Main Area ---
st.header("Infinite Shared 3D World")
st.caption("Move to empty areas to expand the world. Use the sidebar 'Save' controls to store your work.")

# Inject state into JS—including the shared GAME_STATE from our GameState singleton.
injected_state = {
    "ALL_INITIAL_OBJECTS": all_initial_objects,
    "PLOTS_METADATA": plots_metadata,
    "SELECTED_OBJECT_TYPE": st.session_state.selected_object,
    "PLOT_WIDTH": PLOT_WIDTH,
    "PLOT_DEPTH": PLOT_DEPTH,
    "GAME_STATE": game_state.get_state()
}

html_file_path = 'index.html'
html_content_with_state = None

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

    # Use doubled curly braces for literal braces in the injected JS code.
    js_injection_script = f"""
<script>
    window.ALL_INITIAL_OBJECTS = {json.dumps(injected_state["ALL_INITIAL_OBJECTS"])};
    window.PLOTS_METADATA = {json.dumps(injected_state["PLOTS_METADATA"])};
    window.SELECTED_OBJECT_TYPE = {json.dumps(injected_state["SELECTED_OBJECT_TYPE"])};
    window.PLOT_WIDTH = {json.dumps(injected_state["PLOT_WIDTH"])};
    window.PLOT_DEPTH = {json.dumps(injected_state["PLOT_DEPTH"])};
    window.GAME_STATE = {json.dumps(injected_state["GAME_STATE"])};
    console.log("Streamlit State Injected:", {{
        selectedObject: window.SELECTED_OBJECT_TYPE,
        initialObjectsCount: window.ALL_INITIAL_OBJECTS ? window.ALL_INITIAL_OBJECTS.length : 0,
        plotCount: window.PLOTS_METADATA ? window.PLOTS_METADATA.length : 0,
        gameStateObjects: window.GAME_STATE ? window.GAME_STATE.length : 0
    }});
</script>
"""
    html_content_with_state = html_template.replace('</head>', js_injection_script + '\n</head>', 1)
    components.html(html_content_with_state, height=750, scrolling=False)
except FileNotFoundError:
    st.error(f"CRITICAL ERROR: Could not find the file '{html_file_path}'.")
    st.warning(f"Make sure `{html_file_path}` is in the same directory as `app.py` and `{SAVE_DIR}` exists.")
except Exception as e:
    st.error(f"An critical error occurred during HTML preparation or component rendering: {e}")
    st.exception(e)