app.py

# app.py
import streamlit as st
import streamlit.components.v1 as components
import os
import json
import sqlite3 # Use SQLite for robust state management
import uuid
import math
import time # For adding slight delays if needed for debugging FS issues
from streamlit_js_eval import streamlit_js_eval # For JS communication

# --- Constants ---
DB_FILE = "world_state.db" # SQLite database file
PLOT_WIDTH = 50.0
PLOT_DEPTH = 50.0

# --- Database Setup ---
def init_db():
    """Initializes the SQLite database and tables."""
    try:
        # Use WAL mode for potentially better concurrency, though might have visibility delays on some systems
        # conn = sqlite3.connect(DB_FILE, isolation_level=None) # Auto-commit mode
        # cursor.execute('PRAGMA journal_mode=WAL;')
        # conn.close()

        with sqlite3.connect(DB_FILE) as conn: # Use context manager
            cursor = conn.cursor()
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS plots (
                    grid_x INTEGER NOT NULL, grid_z INTEGER NOT NULL, name TEXT,
                    last_updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    PRIMARY KEY (grid_x, grid_z)
                )
            ''')
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS objects (
                    obj_id TEXT PRIMARY KEY, plot_grid_x INTEGER NOT NULL, plot_grid_z INTEGER NOT NULL,
                    type TEXT NOT NULL, pos_x REAL NOT NULL, pos_y REAL NOT NULL, pos_z REAL NOT NULL,
                    rot_x REAL DEFAULT 0.0, rot_y REAL DEFAULT 0.0, rot_z REAL DEFAULT 0.0, rot_order TEXT DEFAULT 'XYZ',
                    FOREIGN KEY (plot_grid_x, plot_grid_z) REFERENCES plots(grid_x, grid_z) ON DELETE CASCADE
                )
            ''')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_objects_plot ON objects (plot_grid_x, plot_grid_z)')
            # No explicit commit needed with 'with' statement unless changes made
        print(f"Database {os.path.abspath(DB_FILE)} initialized/checked.") # Print to console
    except sqlite3.Error as e:
        # Use st.exception for full traceback in the app UI
        st.exception(f"CRITICAL DATABASE INIT ERROR: {e}")
        st.stop()

init_db() # Ensure DB exists and has tables

# --- Helper Functions (Database Operations - Using 'with') ---

# No Caching - always load fresh from DB
def load_world_state_from_db():
    """Loads all plot metadata and object data fresh from the SQLite DB."""
    st.write(f"Executing load_world_state_from_db for session {st.runtime.scriptrunner.get_script_run_ctx().session_id[:5]}...") # Debug session ID
    plots_metadata = []
    all_objects_by_plot = {}
    all_initial_objects_world = []

    try:
        with sqlite3.connect(DB_FILE) as conn: # Auto-closes connection
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()

            st.write("DB Read: Loading plot metadata...")
            cursor.execute("SELECT grid_x, grid_z, name, strftime('%Y-%m-%d %H:%M:%S', last_updated) as updated_ts FROM plots ORDER BY grid_x, grid_z")
            plot_rows = cursor.fetchall()
            st.write(f"DB Read: Found {len(plot_rows)} plot rows.")

            for row in plot_rows:
                gx, gz = row['grid_x'], row['grid_z']
                plot_meta = {
                    'id': f"plot_X{gx}_Z{gz}", 'grid_x': gx, 'grid_z': gz,
                    'name': row['name'] or f"Plot ({gx},{gz})",
                    'x_offset': gx * PLOT_WIDTH, 'z_offset': gz * PLOT_DEPTH,
                    'last_updated': row['updated_ts'] # Add timestamp for debugging
                }
                plots_metadata.append(plot_meta)
                all_objects_by_plot[(gx, gz)] = []

            st.write("DB Read: Loading all objects...")
            cursor.execute("SELECT * FROM objects")
            object_rows = cursor.fetchall()
            st.write(f"DB Read: Found {len(object_rows)} object rows.")

            # Group objects by plot
            for row in object_rows:
                plot_key = (row['plot_grid_x'], row['plot_grid_z'])
                if plot_key in all_objects_by_plot:
                     all_objects_by_plot[plot_key].append(dict(row))

        # Combine and calculate world coordinates (outside DB connection)
        for plot_meta in plots_metadata:
            plot_key = (plot_meta['grid_x'], plot_meta['grid_z'])
            for obj_data in all_objects_by_plot[plot_key]:
                 world_obj_data = obj_data.copy()
                 world_obj_data['pos_x'] = obj_data['pos_x'] + plot_meta['x_offset']
                 world_obj_data['pos_z'] = obj_data['pos_z'] + plot_meta['z_offset']
                 world_obj_data['pos_y'] = obj_data['pos_y']
                 all_initial_objects_world.append(world_obj_data)

        st.write(f"DB Load Complete: {len(plots_metadata)} plots, {len(all_initial_objects_world)} total objects processed.")

    except sqlite3.Error as e:
        st.error(f"Database load error: {e}")
        return [], [] # Return empty on error

    return plots_metadata, all_initial_objects_world


def save_plot_data_to_db(target_grid_x, target_grid_z, objects_data_list):
    """Saves object data list to DB for a specific plot. Overwrites existing objects."""
    st.write(f"Executing save_plot_data_to_db for plot ({target_grid_x},{target_grid_z})...")
    plot_x_offset = target_grid_x * PLOT_WIDTH
    plot_z_offset = target_grid_z * PLOT_DEPTH
    plot_name = f"Plot ({target_grid_x},{target_grid_z})" # Simple default name for upsert

    if not isinstance(objects_data_list, list):
        st.error("Save Error: Invalid object data format (expected list).")
        return False

    try:
        with sqlite3.connect(DB_FILE) as conn: # Auto commit/rollback/close
            cursor = conn.cursor()
            # Use savepoint for finer transaction control within 'with' if needed, but default is fine
            # cursor.execute('BEGIN') # Implicitly handled by 'with' unless isolation_level=None

            # 1. Upsert Plot - ensure plot exists and update timestamp
            st.write(f"DB Save: Upserting plot ({target_grid_x},{target_grid_z})...")
            cursor.execute('''
                INSERT INTO plots (grid_x, grid_z, name) VALUES (?, ?, ?)
                ON CONFLICT(grid_x, grid_z) DO UPDATE SET name = excluded.name, last_updated = CURRENT_TIMESTAMP
            ''', (target_grid_x, target_grid_z, plot_name))
            st.write(f"DB Save: Plot upserted.")


            # 2. Delete ALL existing objects for this specific plot
            st.write(f"DB Save: Deleting old objects for plot ({target_grid_x},{target_grid_z})...")
            cursor.execute("DELETE FROM objects WHERE plot_grid_x = ? AND plot_grid_z = ?", (target_grid_x, target_grid_z))
            st.write(f"DB Save: Deleted {cursor.rowcount} old objects.")

            # 3. Insert the new objects
            insert_count = 0
            objects_to_insert = []
            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 DB save prep: {obj}")
                    continue

                rel_x = pos.get('x', 0.0) - plot_x_offset
                rel_z = pos.get('z', 0.0) - plot_z_offset
                rel_y = pos.get('y', 0.0) # Y is absolute

                objects_to_insert.append((
                    obj_id, target_grid_x, target_grid_z, obj_type,
                    rel_x, rel_y, rel_z,
                    rot.get('_x', 0.0), rot.get('_y', 0.0), rot.get('_z', 0.0),
                    rot.get('_order', 'XYZ')
                ))
                insert_count += 1

            if objects_to_insert:
                st.write(f"DB Save: Inserting {insert_count} new objects...")
                cursor.executemany('''
                    INSERT OR REPLACE INTO objects (obj_id, plot_grid_x, plot_grid_z, type, pos_x, pos_y, pos_z, rot_x, rot_y, rot_z, rot_order)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                ''', objects_to_insert) # Use INSERT OR REPLACE to handle potential obj_id conflicts cleanly
                st.write(f"DB Save: Objects inserted.")
            else:
                 st.write(f"DB Save: No new objects to insert.")

            # 'with' statement handles commit on success, rollback on error

        st.success(f"DB Save Commit: Plot ({target_grid_x},{target_grid_z}) saved with {insert_count} objects.")
        # Add a tiny delay AFTER commit, maybe helps filesystem cache? (Experimental)
        # time.sleep(0.1)
        return True

    except sqlite3.Error as e:
        st.error(f"Database save error for plot ({target_grid_x},{target_grid_z}): {e}")
        st.exception(e) # Show full traceback
        return False


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

# --- Initialize Session State ---
if 'selected_object' not in st.session_state: st.session_state.selected_object = 'None'
if 'js_save_data_result' not in st.session_state: st.session_state.js_save_data_result = None

# --- Load World State From DB (Runs on every script execution/rerun) ---
plots_metadata, all_initial_objects = load_world_state_from_db()

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

    # Refresh Button (Just triggers rerun, load function always hits DB)
    if st.button("🔄 Refresh World View", key="refresh_button"):
        st.info("Reloading world state from database...")
        st.rerun()

    st.header("Navigation (Plots)")
    # ... (Navigation button code remains the same as previous version) ...
    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'] + PLOT_WIDTH / 2
            target_z = plot['z_offset'] + PLOT_DEPTH / 2
            try:
                js_code = f"teleportPlayer({target_x}, {target_z});"
                streamlit_js_eval(js_code=js_code, key=f"teleport_{plot['id']}")
            except Exception as e: st.error(f"Teleport command failed: {e}")
        col_idx = (col_idx + 1) % max_cols

    st.markdown("---")

    # --- Object Placement ---
    st.header("Place Objects")
    # ... (Object selection code remains the same) ...
    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
         # Rerun updates injection, sessionStorage persists JS side

    st.markdown("---")

    # --- Saving ---
    st.header("Save Work")
    st.caption("Saves ALL objects currently within the player's plot to the central database.")
    if st.button("💾 Save Current Plot", key="save_button"):
        js_get_data_code = "getSaveDataAndPosition();"
        streamlit_js_eval(js_code=js_get_data_code, key="js_save_processor")
        st.rerun()


# --- Process Save Data ---
save_data_from_js = st.session_state.get("js_save_processor", None)

if save_data_from_js is not None:
    st.info("Processing save request...")
    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'] # World coords from JS

            if isinstance(objects_to_save, list): # Allow saving empty list (clears plot)
                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)

                st.write(f"Attempting DB save for plot ({target_grid_x},{target_grid_z})")
                # --- Save the data to SQLite DB ---
                save_ok = save_plot_data_to_db(target_grid_x, target_grid_z, objects_to_save)

                if save_ok:
                    # No cache clear needed. Success message is inside save_plot_data_to_db
                    save_processed_successfully = True
                else:
                    st.error(f"Failed DB save for plot ({target_grid_x},{target_grid_z}).")
            else: st.error("Invalid 'objectsToSave' format (expected list).")
        else: st.error("Invalid save payload structure received.")

    except json.JSONDecodeError:
        st.error("Failed to decode save data from client.")
        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)

    # Clear the trigger data from session state ALWAYS after attempting processing
    st.session_state.js_save_processor = None
    # Rerun after processing save to reload world state from DB
    if save_processed_successfully:
        st.rerun()


# --- Main Area ---
st.header("Database Synced 3D World")
st.caption(f"DB: '{os.path.abspath(DB_FILE)}'. Plots loaded: {len(plots_metadata)}. Use 'Refresh' to see updates.")

# --- Load and Prepare HTML ---
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()

    js_injection_script = f"""
<script>
    // Inject data loaded fresh from DB
    window.ALL_INITIAL_OBJECTS = {json.dumps(all_initial_objects)};
    window.PLOTS_METADATA = {json.dumps(plots_metadata)}; // Still useful for JS ground generation
    window.SELECTED_OBJECT_TYPE = {json.dumps(st.session_state.selected_object)};
    window.PLOT_WIDTH = {json.dumps(PLOT_WIDTH)};
    window.PLOT_DEPTH = {json.dumps(PLOT_DEPTH)};
    console.log("Streamlit State Injected:", {{ /* ... logging ... */ }});
</script>
"""
    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, scrolling=False )

except FileNotFoundError:
    st.error(f"CRITICAL ERROR: Could not find the file '{html_file_path}'.")
except Exception as e:
    st.error(f"An critical error occurred during HTML prep/render: {e}")
    st.exception(e)