# app.py
import streamlit as st
import streamlit.components.v1 as components
import os
import json
import pandas as pd
import uuid
from PIL import Image, ImageDraw # Not used for minimap anymore, but kept just in case
from streamlit_js_eval import streamlit_js_eval # For JS communication

# --- Constants ---
SAVE_DIR = "saved_worlds"
PLOT_WIDTH = 50.0 # Width of each plot in 3D space (adjust as needed)
CSV_COLUMNS = ['obj_id', 'type', 'pos_x', 'pos_y', 'pos_z', 'rot_x', 'rot_y', 'rot_z', 'rot_order']

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

# --- Helper Functions ---

@st.cache_data(ttl=3600) # Cache plot list for an hour, or clear manually
def load_plot_metadata():
    """Scans save dir, sorts plots, calculates metadata."""
    plots = []
    # Ensure consistent sorting, e.g., alphabetically which often aligns with plot_001, plot_002 etc.
    try:
        plot_files = sorted([f for f in os.listdir(SAVE_DIR) if f.endswith(".csv")])
    except FileNotFoundError:
        st.error(f"Save directory '{SAVE_DIR}' not found.")
        return [], 0.0 # Return empty if dir doesn't exist
    except Exception as e:
        st.error(f"Error listing save directory '{SAVE_DIR}': {e}")
        return [], 0.0


    current_x_offset = 0.0
    for i, filename in enumerate(plot_files):
        # Extract name - assumes format like 'plot_001_MyName.csv' or just 'plot_001.csv'
        parts = filename[:-4].split('_') # Remove .csv and split by underscore
        plot_id = filename[:-4] # Use filename (without ext) as a unique ID for now
        plot_name = " ".join(parts[1:]) if len(parts) > 1 else f"Plot {i+1}" # Try to extract name

        plots.append({
            'id': plot_id, # Use filename as ID
            'name': plot_name,
            'filename': filename,
            'x_offset': current_x_offset
        })
        current_x_offset += PLOT_WIDTH
    # Also return the offset where the next plot would start
    return plots, current_x_offset

def load_plot_objects(filename, x_offset):
    """Loads objects from a CSV, applying the plot's x_offset."""
    file_path = os.path.join(SAVE_DIR, filename)
    objects = []
    try:
        df = pd.read_csv(file_path)
        # Check if required columns exist, handle gracefully if not
        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 (type, pos_x/y/z). Skipping.")
             return []
        # Ensure optional columns default to something sensible if missing
        # Use vectorized operations for defaults where possible
        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()
            # Apply world offset accumulated during loading
            obj_data['pos_x'] += x_offset
            objects.append(obj_data)
        return objects
    except FileNotFoundError:
        # This shouldn't happen if called via load_plot_metadata results, but handle anyway
        st.error(f"File not found during object load: {filename}")
        return []
    except pd.errors.EmptyDataError:
         # An empty file is valid, represents an empty plot
         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):
    """Saves object data list to a new CSV file, making positions relative."""
    file_path = os.path.join(SAVE_DIR, filename)
    relative_objects = []
    # Ensure objects_data_list is actually a list
    if not isinstance(objects_data_list, list):
        st.error("Invalid data format received for saving (expected a list).")
        print("Invalid save data:", objects_data_list) # Log for debugging
        return False

    for obj in objects_data_list:
        # Validate incoming object structure more carefully
        pos = obj.get('position', {})
        rot = obj.get('rotation', {})
        obj_type = obj.get('type', 'Unknown')
        obj_id = obj.get('obj_id', str(uuid.uuid4())) # Generate ID if missing from JS

        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, # Make relative to plot start
            'pos_y': pos.get('y', 0.0),
            'pos_z': pos.get('z', 0.0),
            '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:
        # Only save if there are objects to save
        if relative_objects:
            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}")
        else:
            # Create an empty file with headers if nothing new was placed
            pd.DataFrame(columns=CSV_COLUMNS).to_csv(file_path, index=False)
            st.info(f"Saved empty plot file: {filename}")
        return True
    except Exception as e:
        st.error(f"Failed to save plot data to {filename}: {e}")
        return False

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

# --- Initialize Session State ---
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 = ""
# Use a more descriptive key for clarity
if 'js_save_data_result' not in st.session_state:
    st.session_state.js_save_data_result = None

# --- Load Plot Metadata ---
# Cached function returns list of plots and the next starting x_offset
plots_metadata, next_plot_x_offset = load_plot_metadata()

# --- Load ALL Objects for Rendering ---
# This could be slow with many plots!
all_initial_objects = []
for plot in plots_metadata:
    all_initial_objects.extend(load_plot_objects(plot['filename'], plot['x_offset']))

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

    st.header("Navigation (Plots)")
    st.caption("Click to teleport player to the start of a plot.")

    # Use columns for a horizontal button layout if desired
    max_cols = 3 # Adjust number of columns
    cols = st.columns(max_cols)
    col_idx = 0
    for plot in plots_metadata:
        # Use an emoji + name for the button
        button_label = f"➡️ {plot.get('name', plot['id'])}" # Fallback to id if name missing
        # Use plot filename (unique) as key
        if cols[col_idx].button(button_label, key=f"nav_{plot['filename']}"):
            # Send command to JS to move the player
            target_x = plot['x_offset']
            try:
                streamlit_js_eval(js_code=f"teleportPlayer({target_x});", key=f"teleport_{plot['filename']}")
            except Exception as e:
                 st.error(f"Failed to send teleport command: {e}")
            # No rerun needed here, JS handles the move instantly

        col_idx = (col_idx + 1) % max_cols # Cycle through columns

    st.markdown("---")

    # --- Object Placement ---
    st.header("Place Objects")
    object_types = ["None", "Simple House", "Tree", "Rock", "Fence Post"]
    current_object_index = 0
    try:
        # Ensure robustness if selected_object is somehow not in list
        current_object_index = object_types.index(st.session_state.selected_object)
    except ValueError:
        st.session_state.selected_object = "None" # Reset to default
        current_object_index = 0

    selected_object_type_widget = st.selectbox(
        "Select Object:",
        options=object_types,
        index=current_object_index,
        key="selected_object_widget" # Use a distinct key for the widget
    )
    # Update session state only if the widget's value actually changes
    # This change WILL trigger a rerun because Streamlit tracks widget state.
    # The JS side now handles preserving its state across the resulting reload via sessionStorage.
    if selected_object_type_widget != st.session_state.selected_object:
         st.session_state.selected_object = selected_object_type_widget
         # We don't *need* to force a rerun here, Streamlit handles it.
         # The important part is that the NEXT run will inject the new selected type,
         # and the JS will restore placed objects from sessionStorage.


    st.markdown("---")

    # --- Saving ---
    st.header("Save New Plot")
    # Ensure text input reflects current state value
    st.session_state.new_plot_name = st.text_input(
        "Name for New Plot:",
        value=st.session_state.new_plot_name,
        placeholder="My Awesome Creation",
        key="new_plot_name_input" # Use distinct key
    )

    if st.button("💾 Save Current Work as New Plot", key="save_button"):
        # 1. Trigger JS function `getSaveData()` defined in index.html
        # This function collects data for newly placed objects and returns a JSON string.
        # The key argument ('js_save_processor') stores the JS result in session_state.
        streamlit_js_eval(
             js_code="getSaveData();",
             key="js_save_processor" # Store result under this key
        )
        # Small delay MAY sometimes help ensure the value is set before rerun, but usually not needed
        # import time
        # time.sleep(0.1)
        st.rerun() # Rerun to process the result in the next step


# --- Process Save Data (if received from JS via the key) ---
# Check the session state key set by the streamlit_js_eval call
save_data_from_js = st.session_state.get("js_save_processor", None)

if save_data_from_js is not None: # Process only if data is present
    st.info("Received save data from client...")
    save_processed_successfully = False
    try:
        # Ensure data is treated as a string before loading json
        if isinstance(save_data_from_js, str):
            objects_to_save = json.loads(save_data_from_js)
        else:
             # Handle case where it might already be parsed by chance (less likely)
             objects_to_save = save_data_from_js

        # Proceed only if we have a list (even an empty one is ok now)
        if isinstance(objects_to_save, list):
            # Determine filename for the new plot
            new_plot_index = len(plots_metadata) # 0-based index -> number of plots
            # Sanitize name: replace spaces, keep only alphanumeric/underscore
            plot_name_sanitized = "".join(c for c in st.session_state.new_plot_name if c.isalnum() or c in (' ')).strip().replace(' ', '_')
            if not plot_name_sanitized: # Ensure there is a name part
                plot_name_sanitized = f"Plot_{new_plot_index + 1}"

            new_filename = f"plot_{new_plot_index:03d}_{plot_name_sanitized}.csv"

            # Save the data, converting world coords to relative coords inside the func
            save_ok = save_plot_data(new_filename, objects_to_save, next_plot_x_offset)

            if save_ok:
                # Clear the plot metadata cache so it reloads with the new file
                load_plot_metadata.clear()
                # Reset the new plot name field for next time
                st.session_state.new_plot_name = ""
                # Reset newly placed objects in JS AFTER successful save
                try:
                     # This call tells JS to clear its internal 'newlyPlacedObjects' array AND sessionStorage
                     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}")

                st.success(f"New plot '{plot_name_sanitized}' saved!")
                save_processed_successfully = True
            else:
                st.error("Failed to save plot data to file.")

        else:
             st.error(f"Received invalid save data format from client (expected list): {type(objects_to_save)}")


    except json.JSONDecodeError:
        st.error("Failed to decode save data from client. Data might be corrupted or empty.")
        print("Received raw data:", save_data_from_js) # Log raw data
    except Exception as e:
        st.error(f"Error processing save: {e}")
        st.exception(e)

    # IMPORTANT: Clear the session state key regardless of success/failure
    # to prevent reprocessing on the next rerun unless the button is clicked again.
    st.session_state.js_save_processor = None

    # Rerun AGAIN after processing save to reflect changes (new plot loaded, cache cleared etc.)
    if save_processed_successfully:
        st.rerun()


# --- Main Area ---
st.header("Shared 3D World")
st.caption("Build side-by-side with others. Saving adds a new plot to the right.")

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

try:
    # --- Read the HTML template file ---
    with open(html_file_path, 'r', encoding='utf-8') as f:
        html_template = f.read()

    # --- Prepare JavaScript code to inject state ---
    # Ensure all state variables are correctly serialized as JSON
    js_injection_script = f"""
<script>
    // Set global variables BEFORE the main script runs
    window.ALL_INITIAL_OBJECTS = {json.dumps(all_initial_objects)}; // All objects from all plots
    window.SELECTED_OBJECT_TYPE = {json.dumps(st.session_state.selected_object)};
    window.PLOT_WIDTH = {json.dumps(PLOT_WIDTH)};
    window.NEXT_PLOT_X_OFFSET = {json.dumps(next_plot_x_offset)}; // Needed for save calculation & ground size
    // Basic logging to verify state in browser console
    console.log("Streamlit State Injected:", {{
        selectedObject: window.SELECTED_OBJECT_TYPE,
        initialObjectsCount: window.ALL_INITIAL_OBJECTS ? window.ALL_INITIAL_OBJECTS.length : 0,
        plotWidth: window.PLOT_WIDTH,
        nextPlotX: window.NEXT_PLOT_X_OFFSET
    }});
</script>
"""
    # --- Inject the script into the HTML template ---
    # Replacing just before </head> is generally safe
    html_content_with_state = html_template.replace('</head>', js_injection_script + '\n</head>', 1)

    # --- Embed HTML Component (ONLY if HTML loading and preparation succeeded) ---
    components.html(
        html_content_with_state,
        height=750, # Adjust height as needed
        scrolling=False
    )

# --- Error Handling ---
except FileNotFoundError:
    st.error(f"CRITICAL 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 critical error occurred during HTML preparation or component rendering: {e}")
    st.exception(e) # Show full traceback for debugging