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 potential delays if needed
from streamlit_js_eval import streamlit_js_eval # For JS communication

# --- Constants ---
DB_FILE = "world_state_v3.db" # Use a new DB file name to ensure fresh start
PLOT_WIDTH = 50.0
PLOT_DEPTH = 50.0

# --- Database Setup ---
def init_db():
    """Initializes the SQLite database and tables."""
    st.write(f"DEBUG: Initializing Database '{os.path.abspath(DB_FILE)}'...")
    try:
        # Ensure connection closes even if errors occur during table creation
        with sqlite3.connect(DB_FILE, timeout=10) as conn: # Added timeout
            cursor = conn.cursor()
            # Enable WAL mode for potentially better concurrency - might help?
            # cursor.execute('PRAGMA journal_mode=WAL;')
            # Plots table
            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)
                )
            ''')
            # Objects table
            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)')
        print(f"DEBUG: Database {DB_FILE} initialized/checked.") # To console
    except sqlite3.Error as e:
        st.exception(f"CRITICAL DATABASE INIT ERROR: {e}")
        st.stop()

init_db() # Run initialization at the start

# --- Helper Functions (Database Operations - Robust Connections) ---

# *** REMOVED @st.cache_data - Load fresh every time ***
def load_world_state_from_db():
    """Loads all plot metadata and object data fresh from the SQLite DB."""
    session_id = st.runtime.scriptrunner.get_script_run_ctx().session_id[:5] # Short ID for logging
    st.write(f"DEBUG [{session_id}]: Executing load_world_state_from_db...")
    plots_metadata = []
    all_initial_objects_world = []

    try:
        # Ensure connection is opened and closed for each load operation
        with sqlite3.connect(DB_FILE, timeout=10) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()

            # 1. Load Plot Metadata
            cursor.execute("SELECT grid_x, grid_z, name FROM plots ORDER BY grid_x, grid_z")
            plot_rows = cursor.fetchall()
            st.write(f"DEBUG [{session_id}]: DB Read Found {len(plot_rows)} plot rows.")
            if not plot_rows:
                 st.write(f"DEBUG [{session_id}]: No plots found in DB.")

            plot_keys = set() # To efficiently map objects later
            for row in plot_rows:
                gx, gz = row['grid_x'], row['grid_z']
                plots_metadata.append({
                    '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
                })
                plot_keys.add((gx, gz))

            # 2. Load All Objects
            cursor.execute("SELECT obj_id, plot_grid_x, plot_grid_z, type, pos_x, pos_y, pos_z, rot_x, rot_y, rot_z, rot_order FROM objects")
            object_rows = cursor.fetchall()
            st.write(f"DEBUG [{session_id}]: DB Read Found {len(object_rows)} object rows.")

            # 3. Combine and Calculate World Coordinates
            objects_loaded_count = 0
            for row in object_rows:
                plot_key = (row['plot_grid_x'], row['plot_grid_z'])
                # Ensure the object belongs to a known plot (data integrity check)
                if plot_key in plot_keys:
                    obj_data = dict(row) # Convert row object to dict
                    # Find corresponding plot metadata for offset calculation
                    plot_meta = next((p for p in plots_metadata if p['grid_x'] == row['plot_grid_x'] and p['grid_z'] == row['plot_grid_z']), None)
                    if plot_meta:
                        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'] # Y is already world Y
                        all_initial_objects_world.append(world_obj_data)
                        objects_loaded_count += 1
                    else:
                         print(f"WARN: Object {obj_data['obj_id']} references non-existent plot {plot_key}") # To console
                else:
                     print(f"WARN: Object {obj_data['obj_id']} references unknown plot {plot_key}") # To console


        st.write(f"DEBUG [{session_id}]: DB Load Complete. Processed {len(plots_metadata)} plots, {objects_loaded_count} objects.")

    except sqlite3.Error as e:
        st.exception(f"Database load error for session {session_id}: {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"DEBUG: 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
    # Use a generic name or allow passing one if needed later
    plot_name = f"Plot ({target_grid_x},{target_grid_z})"

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

    try:
        # Use 'with' for automatic transaction handling and closing
        with sqlite3.connect(DB_FILE, timeout=10) as conn:
            cursor = conn.cursor()

            # 1. Upsert Plot (ensures plot exists, updates timestamp)
            cursor.execute('''
                INSERT INTO plots (grid_x, grid_z, name) VALUES (?, ?, ?)
                ON CONFLICT(grid_x, grid_z) DO UPDATE SET last_updated = CURRENT_TIMESTAMP
            ''', (target_grid_x, target_grid_z, plot_name))
            # st.write(f"DEBUG: Upserted plot metadata for ({target_grid_x},{target_grid_z})")

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

            # 3. Insert the new objects
            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': 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)

                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')
                ))

            if objects_to_insert:
                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 based on obj_id primary key
                st.write(f"DEBUG: Inserted/Replaced {len(objects_to_insert)} objects.")
            else:
                 st.write(f"DEBUG: No objects provided to insert for plot ({target_grid_x},{target_grid_z}). Plot cleared.")

            # Commit happens automatically when 'with' block exits without error

        st.success(f"Plot ({target_grid_x},{target_grid_z}) saved successfully to DB.")
        return True

    except sqlite3.Error as e:
        st.exception(f"DATABASE SAVE ERROR for plot ({target_grid_x},{target_grid_z}): {e}")
        return False

# --- Page Config ---
st.set_page_config( page_title="DB Synced World Builder v3", 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
# No refresh counter needed

# --- Load World State From DB (Fresh on each run/rerun) ---
plots_metadata, all_initial_objects = load_world_state_from_db()

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

    # Refresh Button (Just reruns the script)
    if st.button("🔄 Refresh World View", key="refresh_button"):
        st.info("Reloading world state from database...")
        st.rerun() # Rerun forces call to load_world_state_from_db

    st.header("Navigation (Plots)")
    # ... (Navigation button code unchanged - uses latest plots_metadata) ...
    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: streamlit_js_eval(js_code=f"teleportPlayer({target_x}, {target_z});", 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 unchanged) ...
    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("---")

    # --- Saving ---
    st.header("Save Work")
    st.caption("Saves ALL objects in the player's current 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.write("DEBUG: Processing save request from JS...") # Debug
    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):
                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)
                # --- 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: save_processed_successfully = True # Flag success
                else: st.error(f"Failed DB save for plot ({target_grid_x},{target_grid_z}).")
            else: st.error("Save Error: 'objectsToSave' format invalid (expected list).")
        else: st.error("Save Error: Invalid payload structure received.")
    except json.JSONDecodeError: st.error("Save Error: Failed to decode JSON data from client.")
    except Exception as e: st.exception(f"Save Error: An unexpected error occurred: {e}")

    # Clear the trigger data from session state ALWAYS after processing attempt
    st.session_state.js_save_processor = None
    # Rerun ONLY if save was flagged successful to reload the state
    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' button to sync.")

# --- 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()

    # Inject data loaded fresh from DB
    js_injection_script = f"""
<script>
    window.ALL_INITIAL_OBJECTS = {json.dumps(all_initial_objects)};
    window.PLOTS_METADATA = {json.dumps(plots_metadata)};
    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:", {{ /* Basic 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 file '{html_file_path}'.")
except Exception as e: st.exception(f"HTML prep/render error: {e}")