app.py

# app.py (Refactored & Consolidated - Checked Indentation)
import streamlit as st
import asyncio
import websockets
import uuid
from datetime import datetime
import os
import random
import time
import hashlib
import glob
import base64
import io
import streamlit.components.v1 as components
import edge_tts
import nest_asyncio
import re
import pytz
import shutil
from PyPDF2 import PdfReader
import threading
import json
import zipfile
from dotenv import load_dotenv
from streamlit_marquee import streamlit_marquee
from collections import defaultdict, Counter
import pandas as pd # Still used for fallback CSV load? Keep for now.
from streamlit_js_eval import streamlit_js_eval
from PIL import Image # Needed for paste_image_component

# ==============================================================================
# Configuration & Constants
# ==============================================================================

# 🛠️ Patch asyncio for nesting
nest_asyncio.apply()

# 🎨 Page Config
st.set_page_config(
    page_title="🤖🏗️ Shared World Builder 🏆",
    page_icon="🏗️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# General Constants
icons = '🤖🏗️🗣️💾'
Site_Name = '🤖🏗️ Shared World Builder 🗣️'
START_ROOM = "World Lobby 🌍"
MEDIA_DIR = "." # Base directory for general files
STATE_FILE = "user_state.txt" # For remembering username

# User/Chat Constants
FUN_USERNAMES = {
    "BuilderBot 🤖": "en-US-AriaNeural", "WorldWeaver 🕸️": "en-US-JennyNeural",
    "Terraformer 🌱": "en-GB-SoniaNeural", "SkyArchitect ☁️": "en-AU-NatashaNeural",
    "PixelPainter 🎨": "en-CA-ClaraNeural", "VoxelVortex 🌪️": "en-US-GuyNeural",
    "CosmicCrafter ✨": "en-GB-RyanNeural", "GeoGuru 🗺️": "en-AU-WilliamNeural",
    "BlockBard 🧱": "en-CA-LiamNeural", "SoundSculptor 🔊": "en-US-AnaNeural",
}
EDGE_TTS_VOICES = list(set(FUN_USERNAMES.values()))
CHAT_DIR = "chat_logs"

# Audio Constants
AUDIO_CACHE_DIR = "audio_cache"
AUDIO_DIR = "audio_logs"

# World Builder Constants
SAVED_WORLDS_DIR = "saved_worlds" # Directory for MD world files
PLOT_WIDTH = 50.0 # Needed for JS injection
PLOT_DEPTH = 50.0 # Needed for JS injection
WORLD_STATE_FILE_MD_PREFIX = "🌍_" # Prefix for world save files

# File Emojis
FILE_EMOJIS = {"md": "📝", "mp3": "🎵", "png": "🖼️", "mp4": "🎥", "zip": "📦", "json": "📄"}


# --- Directories ---
for d in [CHAT_DIR, AUDIO_DIR, AUDIO_CACHE_DIR, SAVED_WORLDS_DIR]:
    os.makedirs(d, exist_ok=True)

# --- API Keys (Placeholder) ---
load_dotenv()
# ANTHROPIC_KEY = os.getenv('ANTHROPIC_API_KEY', st.secrets.get('ANTHROPIC_API_KEY', ""))
# OPENAI_KEY = os.getenv('OPENAI_API_KEY', st.secrets.get('OPENAI_API_KEY', ""))

# --- Global State & Locks ---
world_objects_lock = threading.Lock()
world_objects = defaultdict(dict) # In-memory world state {obj_id: data}
connected_clients = set() # Holds client_id strings

# ==============================================================================
# Utility Functions
# ==============================================================================

def get_current_time_str(tz='UTC'):
    """Gets formatted timestamp string in specified timezone (default UTC)."""
    try:
        timezone = pytz.timezone(tz)
        now_aware = datetime.now(timezone)
    except pytz.UnknownTimeZoneError:
        now_aware = datetime.now(pytz.utc)
    except Exception as e:
         print(f"Timezone error ({tz}), using UTC. Error: {e}")
         now_aware = datetime.now(pytz.utc)
    return now_aware.strftime('%Y%m%d_%H%M%S')


def clean_filename_part(text, max_len=30):
    """Cleans a string part for use in a filename."""
    if not isinstance(text, str): text = "invalid_name"
    text = re.sub(r'\s+', '_', text) # Replace spaces
    text = re.sub(r'[^\w\-.]', '', text) # Keep word chars, hyphen, period
    return text[:max_len]

def run_async(async_func, *args, **kwargs):
    """Runs an async function safely from a sync context using create_task."""
    # This helper attempts to schedule the async function as a background task
    # without blocking the main Streamlit thread.
    try:
        loop = asyncio.get_running_loop()
        return loop.create_task(async_func(*args, **kwargs))
    except RuntimeError: # No running loop in this thread
        # Fallback: Run in a new loop (might block slightly, less ideal for UI responsiveness)
        # Consider if truly background execution is needed (e.g., ThreadPoolExecutor)
        print(f"Warning: Running async func {async_func.__name__} in new event loop.")
        try:
            return asyncio.run(async_func(*args, **kwargs))
        except Exception as e:
            print(f"Error running async func {async_func.__name__} in new loop: {e}")
            return None
    except Exception as e:
        print(f"Error scheduling async task {async_func.__name__}: {e}")
        return None

def ensure_dir(dir_path):
     """Creates directory if it doesn't exist."""
     os.makedirs(dir_path, exist_ok=True)

# Ensure directories exist on startup
for d in [CHAT_DIR, AUDIO_DIR, AUDIO_CACHE_DIR, SAVED_WORLDS_DIR]:
    ensure_dir(d)

# ==============================================================================
# World State File Handling (Markdown + JSON)
# ==============================================================================

def generate_world_save_filename(name="World"):
    """Generates a filename for saving world state MD files."""
    timestamp = get_current_time_str() # Use UTC for consistency
    clean_name = clean_filename_part(name)
    rand_hash = hashlib.md5(str(time.time()).encode() + name.encode()).hexdigest()[:6] # Seed hash
    return f"{WORLD_STATE_FILE_MD_PREFIX}{clean_name}_{timestamp}_{rand_hash}.md"

def parse_world_filename(filename):
    """Extracts info from filename if possible, otherwise returns defaults."""
    basename = os.path.basename(filename)
    # Check prefix and suffix
    if basename.startswith(WORLD_STATE_FILE_MD_PREFIX) and basename.endswith(".md"):
        # Remove prefix and suffix before splitting
        core_name = basename[len(WORLD_STATE_FILE_MD_PREFIX):-3]
        parts = core_name.split('_')
        if len(parts) >= 3: # Expecting Name_Timestamp_Hash
            timestamp_str = parts[-2]
            # Combine parts before timestamp and hash for the name
            name_parts = parts[:-2]
            name = "_".join(name_parts) if name_parts else "Untitled" # Handle empty name parts
            dt_obj = None
            try: # Try parsing timestamp
                dt_obj = datetime.strptime(timestamp_str, '%Y%m%d_%H%M%S')
                dt_obj = pytz.utc.localize(dt_obj) # Assume UTC
            except (ValueError, pytz.exceptions.AmbiguousTimeError, pytz.exceptions.NonExistentTimeError):
                dt_obj = None # Parsing failed
            return {"name": name.replace('_', ' '), "timestamp": timestamp_str, "dt": dt_obj, "filename": filename}

    # Fallback for unknown format
    dt_fallback = None
    try:
         mtime = os.path.getmtime(filename)
         dt_fallback = datetime.fromtimestamp(mtime, tz=pytz.utc)
    except Exception: pass
    return {"name": basename.replace('.md',''), "timestamp": "Unknown", "dt": dt_fallback, "filename": filename}


def save_world_state_to_md(target_filename_base):
    """Saves the current in-memory world state to a specific MD file (basename)."""
    global world_objects
    save_path = os.path.join(SAVED_WORLDS_DIR, target_filename_base)
    print(f"Acquiring lock to save world state to: {save_path}...")
    success = False
    with world_objects_lock:
        # Create a deep copy for saving if needed, dict() might be shallow
        world_data_dict = dict(world_objects) # Convert defaultdict for saving
        print(f"Saving {len(world_data_dict)} objects...")
        # Use the target filename to generate header info
        parsed_info = parse_world_filename(save_path) # Parse the full path/intended name
        timestamp_save = get_current_time_str()
        md_content = f"""# World State: {parsed_info['name']}
* **File Saved:** {timestamp_save} (UTC)
* **Source Timestamp:** {parsed_info['timestamp']}
* **Objects:** {len(world_data_dict)}

```json
{json.dumps(world_data_dict, indent=2)}
```"""
        try:
            ensure_dir(SAVED_WORLDS_DIR)
            with open(save_path, 'w', encoding='utf-8') as f: f.write(md_content)
            print(f"World state saved successfully to {target_filename_base}")
            success = True
        except Exception as e:
            print(f"Error saving world state to {save_path}: {e}")
            # Avoid st.error in potentially non-main thread
    return success


def load_world_state_from_md(filename_base):
    """Loads world state from an MD file (basename), updates global state, returns success bool."""
    global world_objects
    load_path = os.path.join(SAVED_WORLDS_DIR, filename_base)
    print(f"Loading world state from MD file: {load_path}...")
    if not os.path.exists(load_path):
        st.error(f"World file not found: {filename_base}")
        return False

    try:
        with open(load_path, 'r', encoding='utf-8') as f: content = f.read()
        # More robust JSON extraction
        json_match = re.search(r"```json\s*(\{[\s\S]*?\})\s*```", content, re.IGNORECASE)
        if not json_match: st.error(f"Could not find valid JSON block in {filename_base}"); return False

        world_data_dict = json.loads(json_match.group(1))

        print(f"Acquiring lock to update world state from {filename_base}...")
        with world_objects_lock:
             world_objects.clear()
             for k, v in world_data_dict.items(): world_objects[str(k)] = v
             loaded_count = len(world_objects)
        print(f"Loaded {loaded_count} objects from {filename_base}. Lock released.")
        st.session_state.current_world_file = filename_base # Track loaded file (basename)
        return True

    except json.JSONDecodeError as e: st.error(f"Invalid JSON found in {filename_base}: {e}"); return False
    except Exception as e: st.error(f"Error loading world state from {filename_base}: {e}"); st.exception(e); return False

def get_saved_worlds():
    """Scans the saved worlds directory for world MD files and parses them."""
    try:
        ensure_dir(SAVED_WORLDS_DIR)
        # Use the prefix in the glob pattern
        world_files = glob.glob(os.path.join(SAVED_WORLDS_DIR, f"{WORLD_STATE_FILE_MD_PREFIX}*.md"))
        parsed_worlds = [parse_world_filename(f) for f in world_files]
        # Sort by datetime object (newest first), handle None dt values
        parsed_worlds.sort(key=lambda x: x['dt'] if x['dt'] else datetime.min.replace(tzinfo=pytz.utc), reverse=True)
        return parsed_worlds
    except Exception as e:
         print(f"Error scanning saved worlds: {e}")
         st.error(f"Could not scan saved worlds: {e}")
         return []

# ==============================================================================
# User State & Session Init
# ==============================================================================

def save_username(username):
    try:
        with open(STATE_FILE, 'w') as f: f.write(username)
    except Exception as e: print(f"Failed save username: {e}")

def load_username():
    if os.path.exists(STATE_FILE):
        try:
            with open(STATE_FILE, 'r') as f: return f.read().strip()
        except Exception as e: print(f"Failed load username: {e}")
    return None

def init_session_state():
    """Initializes Streamlit session state variables."""
    defaults = {
        'server_running_flag': False, 'server_instance': None, 'server_task': None,
        'active_connections': defaultdict(dict), 'last_chat_update': 0, 'message_input': "",
        'audio_cache': {}, 'tts_voice': "en-US-AriaNeural", 'chat_history': [],
        'marquee_settings': {"background": "#1E1E1E", "color": "#FFFFFF", "font-size": "14px", "animationDuration": "20s", "width": "100%", "lineHeight": "35px"},
        'enable_audio': True, 'download_link_cache': {}, 'username': None, 'autosend': False,
        'last_message': "", 'timer_start': time.time(), 'last_sent_transcript': "",
        'last_refresh': time.time(), 'auto_refresh': False, 'refresh_rate': 30,
        'selected_object': 'None', 'initial_world_state_loaded': False,
        'current_world_file': None, # Track loaded world filename (basename)
        'operation_timings': {}, 'performance_metrics': defaultdict(list),
        'paste_image_base64': "", 'new_world_name': "MyWorld"
    }
    for k, v in defaults.items():
        if k not in st.session_state: st.session_state[k] = v
    # Ensure complex types initialized correctly
    if not isinstance(st.session_state.active_connections, defaultdict): st.session_state.active_connections = defaultdict(dict)
    if not isinstance(st.session_state.chat_history, list): st.session_state.chat_history = []
    if not isinstance(st.session_state.marquee_settings, dict): st.session_state.marquee_settings = defaults['marquee_settings']
    if not isinstance(st.session_state.audio_cache, dict): st.session_state.audio_cache = {}
    if not isinstance(st.session_state.download_link_cache, dict): st.session_state.download_link_cache = {}

# ==============================================================================
# Audio / TTS / Chat / File Handling Helpers
# ==============================================================================

# --- Text & File Helpers ---
def clean_text_for_tts(text):
    if not isinstance(text, str): return "No text"
    text = re.sub(r'\[([^\]]+)\]\([^\)]+\)', r'\1', text)
    text = re.sub(r'[#*_`!]', '', text)
    text = ' '.join(text.split())
    return text[:250] or "No text"

def create_file(content, username, file_type="md", save_path=None):
    if not save_path:
         filename = generate_filename(content, username, file_type)
         save_path = os.path.join(MEDIA_DIR, filename)
    ensure_dir(os.path.dirname(save_path))
    try:
        with open(save_path, 'w', encoding='utf-8') as f: f.write(content)
        print(f"Created file: {save_path}"); return save_path
    except Exception as e: print(f"Error creating file {save_path}: {e}"); return None

def get_download_link(file_path, file_type="md"):
    """Generates a base64 download link for a given file."""
    if not file_path or not os.path.exists(file_path):
         basename = os.path.basename(file_path) if file_path else "N/A"
         return f"<small>Not found: {basename}</small>"
    try: mtime = os.path.getmtime(file_path)
    except OSError: mtime = 0
    cache_key = f"dl_{file_path}_{mtime}"
    if 'download_link_cache' not in st.session_state: st.session_state.download_link_cache = {}
    if cache_key not in st.session_state.download_link_cache:
        try:
            with open(file_path, "rb") as f: b64 = base64.b64encode(f.read()).decode()
            mime_types = {"md": "text/markdown", "mp3": "audio/mpeg", "png": "image/png", "mp4": "video/mp4", "zip": "application/zip", "json": "application/json"}
            basename = os.path.basename(file_path)
            # Changed emoji and text for clarity
            link_html = f'<a href="data:{mime_types.get(file_type, "application/octet-stream")};base64,{b64}" download="{basename}" title="Download {basename}">{FILE_EMOJIS.get(file_type, "📄")}</a>'
            st.session_state.download_link_cache[cache_key] = link_html
        except Exception as e:
             print(f"Error generating DL link for {file_path}: {e}")
             return f"<small>Err</small>"
    return st.session_state.download_link_cache.get(cache_key, "<small>CacheErr</small>")

# --- Audio / TTS ---
async def async_edge_tts_generate(text, voice, username):
    """Generates TTS audio using EdgeTTS and caches the result."""
    if not text: return None
    cache_key = hashlib.md5(f"{text[:150]}_{voice}".encode()).hexdigest()
    if 'audio_cache' not in st.session_state: st.session_state.audio_cache = {}
    cached_path = st.session_state.audio_cache.get(cache_key)
    if cached_path and os.path.exists(cached_path): return cached_path

    text_cleaned = clean_text_for_tts(text);
    if not text_cleaned or text_cleaned == "No text": return None
    filename_base = generate_filename(text_cleaned, username, "mp3"); save_path = os.path.join(AUDIO_DIR, filename_base);
    ensure_dir(AUDIO_DIR)
    try:
        communicate = edge_tts.Communicate(text_cleaned, voice); await communicate.save(save_path);
        if os.path.exists(save_path) and os.path.getsize(save_path) > 0:
             st.session_state.audio_cache[cache_key] = save_path; return save_path
        else: print(f"Audio file {save_path} failed generation."); return None
    except Exception as e: print(f"Edge TTS Error: {e}"); return None

def play_and_download_audio(file_path):
    """Displays audio player and download link in Streamlit."""
    if file_path and os.path.exists(file_path):
        try:
             st.audio(file_path)
             file_type = file_path.split('.')[-1]
             st.markdown(get_download_link(file_path, file_type), unsafe_allow_html=True)
        except Exception as e: st.error(f"Audio display error for {os.path.basename(file_path)}: {e}")
    # else: st.warning(f"Audio file not found: {os.path.basename(file_path) if file_path else 'N/A'}") # Can be noisy

# --- Chat ---
async def save_chat_entry(username, message, voice, is_markdown=False):
    """Saves chat entry to a file and session state, generates audio."""
    if not message.strip(): return None, None
    timestamp_str = get_current_time_str();
    entry = f"[{timestamp_str}] {username} ({voice}): {message}" if not is_markdown else f"[{timestamp_str}] {username} ({voice}):\n```markdown\n{message}\n```"
    md_filename_base = generate_filename(message, username, "md"); md_file_path = os.path.join(CHAT_DIR, md_filename_base);
    md_file = create_file(entry, username, "md", save_path=md_file_path) # Save to file
    if 'chat_history' not in st.session_state: st.session_state.chat_history = [];
    st.session_state.chat_history.append(entry) # Add to live history
    audio_file = None;
    if st.session_state.get('enable_audio', True):
         tts_message = message # Use original message for TTS
         audio_file = await async_edge_tts_generate(tts_message, voice, username)
    return md_file, audio_file

async def load_chat_history():
    """Loads chat history from files if session state is empty."""
    # This ensures history is loaded once per session if needed
    if 'chat_history' not in st.session_state: st.session_state.chat_history = []
    if not st.session_state.chat_history:
        ensure_dir(CHAT_DIR)
        print("Loading chat history from files...")
        chat_files = sorted(glob.glob(os.path.join(CHAT_DIR, "*.md")), key=os.path.getmtime); loaded_count = 0
        temp_history = []
        for f_path in chat_files:
             try:
                  with open(f_path, 'r', encoding='utf-8') as file: temp_history.append(file.read().strip()); loaded_count += 1
             except Exception as e: print(f"Err read chat {f_path}: {e}")
        st.session_state.chat_history = temp_history # Assign loaded history
        print(f"Loaded {loaded_count} chat entries from files.")
    return st.session_state.chat_history

# --- File Management ---
def create_zip_of_files(files_to_zip, prefix="Archive"):
    if not files_to_zip: st.warning("No files provided to zip."); return None
    timestamp = format_timestamp_prefix(f"Zip_{prefix}"); zip_name = f"{prefix}_{timestamp}.zip"
    try:
        print(f"Creating zip: {zip_name}...");
        with zipfile.ZipFile(zip_name, 'w', zipfile.ZIP_DEFLATED) as z:
            for f in files_to_zip:
                 if os.path.exists(f): z.write(f, os.path.basename(f))
                 else: print(f"Skip zip missing: {f}")
        print("Zip success."); st.success(f"Created {zip_name}"); return zip_name
    except Exception as e: print(f"Zip failed: {e}"); st.error(f"Zip failed: {e}"); return None

def delete_files(file_patterns, exclude_files=None):
    """Deletes files matching patterns, excluding protected/specified files."""
    # Core protected files
    protected = [STATE_FILE, "app.py", "index.html", "requirements.txt", "README.md"]
    # Dynamically protect currently loaded world file if specified
    current_world = st.session_state.get('current_world_file')
    if current_world: protected.append(current_world)
    # Add user exclusions
    if exclude_files: protected.extend(exclude_files)

    deleted_count = 0; errors = 0
    for pattern in file_patterns:
         # Expand pattern relative to current directory or specified dir
         pattern_path = pattern # Assume pattern includes path if needed (e.g., from os.path.join)
         print(f"Attempting to delete files matching: {pattern_path}")
         try:
              files_to_delete = glob.glob(pattern_path)
              if not files_to_delete: print(f"No files found for pattern: {pattern}"); continue

              for f_path in files_to_delete:
                   basename = os.path.basename(f_path)
                   # Check if it's a file and NOT protected
                   if os.path.isfile(f_path) and basename not in protected:
                        try: os.remove(f_path); print(f"Deleted: {f_path}"); deleted_count += 1
                        except Exception as e: print(f"Failed delete {f_path}: {e}"); errors += 1
                   elif os.path.isdir(f_path): print(f"Skipping directory: {f_path}")
                   #else: print(f"Skipping protected/non-file: {f_path}") # Debugging
         except Exception as glob_e: print(f"Error matching pattern {pattern}: {glob_e}"); errors += 1

    msg = f"Deleted {deleted_count} files.";
    if errors > 0: msg += f" Encountered {errors} errors."; st.warning(msg)
    elif deleted_count > 0: st.success(msg)
    else: st.info("No matching files found to delete.")
    # Clear relevant caches
    st.session_state['download_link_cache'] = {}
    st.session_state['audio_cache'] = {}


# --- Image Handling ---
async def save_pasted_image(image, username):
     if not image: return None
     try:
          img_hash = hashlib.md5(image.tobytes()).hexdigest()[:8]; timestamp = format_timestamp_prefix(username); filename = f"{timestamp}_pasted_{img_hash}.png"; filepath = os.path.join(MEDIA_DIR, filename)
          image.save(filepath, "PNG"); print(f"Pasted image saved: {filepath}"); return filepath
     except Exception as e: print(f"Failed image save: {e}"); return None

def paste_image_component():
      pasted_img = None; img_type = None
      with st.form(key="paste_form"):
           paste_input = st.text_area("Paste Image Data Here", key="paste_input_area", height=50); submit_button = st.form_submit_button("Paste Image 📋")
           if submit_button and paste_input and paste_input.startswith('data:image'):
               try:
                   mime_type = paste_input.split(';')[0].split(':')[1]; base64_str = paste_input.split(',')[1]; img_bytes = base64.b64decode(base64_str); pasted_img = Image.open(io.BytesIO(img_bytes)); img_type = mime_type.split('/')[1]
                   st.image(pasted_img, caption=f"Pasted ({img_type.upper()})", width=150); st.session_state.paste_image_base64 = base64_str
               except ImportError: st.error("Pillow library needed for image pasting.")
               except Exception as e: st.error(f"Img decode err: {e}"); st.session_state.paste_image_base64 = ""
           elif submit_button: st.warning("No valid img data."); st.session_state.paste_image_base64 = ""
      # Return the image object if successfully pasted and submitted in THIS RUN
      return pasted_img if submit_button and pasted_img else None


# --- PDF Processing ---
class AudioProcessor:
      def __init__(self): self.cache_dir=AUDIO_CACHE_DIR; ensure_dir(self.cache_dir); self.metadata=json.load(open(f"{self.cache_dir}/metadata.json", 'r')) if os.path.exists(f"{self.cache_dir}/metadata.json") else {}
      def _save_metadata(self):
            try:
                 with open(f"{self.cache_dir}/metadata.json", 'w') as f: json.dump(self.metadata, f, indent=2)
            except Exception as e: print(f"Failed metadata save: {e}")
      async def create_audio(self, text, voice='en-US-AriaNeural'):
            cache_key=hashlib.md5(f"{text[:150]}:{voice}".encode()).hexdigest(); cache_path=os.path.join(self.cache_dir, f"{cache_key}.mp3") # Use join
            if cache_key in self.metadata and os.path.exists(cache_path): return cache_path
            text_cleaned=clean_text_for_tts(text);
            if not text_cleaned: return None
            ensure_dir(os.path.dirname(cache_path))
            try:
                 communicate=edge_tts.Communicate(text_cleaned,voice); await communicate.save(cache_path)
                 if os.path.exists(cache_path) and os.path.getsize(cache_path) > 0:
                    self.metadata[cache_key]={'timestamp': datetime.now().isoformat(), 'text_length': len(text_cleaned), 'voice': voice}; self._save_metadata()
                    return cache_path
                 else: return None
            except Exception as e: print(f"TTS Create Audio Error: {e}"); return None

def process_pdf_tab(pdf_file, max_pages, voice):
      st.subheader("PDF Processing")
      if pdf_file is None: st.info("Upload a PDF file to begin."); return
      audio_processor = AudioProcessor()
      try:
          reader=PdfReader(pdf_file)
          if reader.is_encrypted: st.warning("PDF is encrypted."); return
          total_pages=min(len(reader.pages),max_pages);
          st.write(f"Processing first {total_pages} pages of '{pdf_file.name}'...");
          texts, audios={}, {}; page_threads = []; results_lock = threading.Lock()

          def process_page_sync(page_num, page_text):
              async def run_async_audio(): return await audio_processor.create_audio(page_text, voice)
              try:
                   # Use the run_async helper
                   audio_path = run_async(run_async_audio).result() # Blocking wait here might be okay for thread
                   if audio_path:
                        with results_lock: audios[page_num] = audio_path
              except Exception as page_e: print(f"Err process page {page_num+1}: {page_e}")

          # Start threads
          for i in range(total_pages):
               try:
                   page = reader.pages[i]; text = page.extract_text();
                   if text and text.strip():
                       texts[i]=text; thread = threading.Thread(target=process_page_sync, args=(i, text)); page_threads.append(thread); thread.start()
                   else: texts[i] = "[No text extracted]"
               except Exception as extract_e: texts[i] = f"[Error extract: {extract_e}]"; print(f"Error page {i+1} extract: {extract_e}")

          # Wait for threads and display progress
          progress_bar = st.progress(0.0)
          total_threads = len(page_threads)
          start_join_time = time.time()
          while any(t.is_alive() for t in page_threads):
              completed_threads = total_threads - sum(t.is_alive() for t in page_threads)
              progress = completed_threads / total_threads if total_threads > 0 else 1.0
              progress_bar.progress(min(progress, 1.0)) # Cap at 1.0
              if time.time() - start_join_time > 600: # Timeout after 10 mins
                   print("PDF processing timed out waiting for threads.")
                   break
              time.sleep(0.5)
          progress_bar.progress(1.0)

          # Display results
          st.write("Processing complete. Displaying results:")
          for i in range(total_pages):
                with st.expander(f"Page {i+1}"):
                     st.markdown(texts.get(i, "[Error getting text]"))
                     audio_file = audios.get(i)
                     if audio_file: play_and_download_audio(audio_file)
                     else: st.caption("Audio generation failed or was skipped.")

      except Exception as pdf_e: st.error(f"Err read PDF: {pdf_e}"); st.exception(pdf_e)


# ==============================================================================
# WebSocket Server Logic
# ==============================================================================

async def register_client(websocket):
    client_id = str(websocket.id); connected_clients.add(client_id);
    if 'active_connections' not in st.session_state: st.session_state.active_connections = defaultdict(dict);
    st.session_state.active_connections[client_id] = websocket; print(f"Client registered: {client_id}. Total: {len(connected_clients)}")

async def unregister_client(websocket):
    client_id = str(websocket.id); connected_clients.discard(client_id);
    if 'active_connections' in st.session_state: st.session_state.active_connections.pop(client_id, None);
    print(f"Client unregistered: {client_id}. Remaining: {len(connected_clients)}")

async def send_safely(websocket, message, client_id):
     """Wrapper to send message and handle potential connection errors."""
     try: await websocket.send(message)
     except websockets.ConnectionClosed: print(f"WS Send failed (Closed) client {client_id}"); raise
     except RuntimeError as e: print(f"WS Send failed (Runtime {e}) client {client_id}"); raise
     except Exception as e: print(f"WS Send failed (Other {e}) client {client_id}"); raise

async def broadcast_message(message, exclude_id=None):
    """Sends a message to all connected clients except the excluded one."""
    if not connected_clients: return
    tasks = []; current_client_ids = list(connected_clients); active_connections_copy = st.session_state.active_connections.copy()
    for client_id in current_client_ids:
        if client_id == exclude_id: continue
        websocket = active_connections_copy.get(client_id)
        if websocket: tasks.append(asyncio.create_task(send_safely(websocket, message, client_id)))
    if tasks: await asyncio.gather(*tasks, return_exceptions=True) # Gather results/exceptions

async def broadcast_world_update():
    """Broadcasts the current world state to all clients."""
    with world_objects_lock: current_state_payload = dict(world_objects)
    update_msg = json.dumps({"type": "initial_state", "payload": current_state_payload}) # Force full refresh
    print(f"Broadcasting full world update ({len(current_state_payload)} objects)...")
    await broadcast_message(update_msg)

async def websocket_handler(websocket, path):
    """Handles WebSocket connections and messages."""
    await register_client(websocket); client_id = str(websocket.id);
    # Use username from main session state - ASSUMES session state is accessible here.
    # This might be unreliable depending on how threads/asyncio interact with Streamlit's context.
    # A safer approach might involve passing necessary user info during registration if needed.
    username = st.session_state.get('username', f"User_{client_id[:4]}")

    try: # Send initial state
        with world_objects_lock: initial_state_payload = dict(world_objects)
        initial_state_msg = json.dumps({"type": "initial_state", "payload": initial_state_payload}); await websocket.send(initial_state_msg)
        print(f"Sent initial state ({len(initial_state_payload)} objs) to {client_id}")
        await broadcast_message(json.dumps({"type": "user_join", "payload": {"username": username, "id": client_id}}), exclude_id=client_id)
    except Exception as e: print(f"Error initial phase {client_id}: {e}")

    try: # Message loop
        async for message in websocket:
            try:
                data = json.loads(message); msg_type = data.get("type"); payload = data.get("payload", {});
                sender_username = payload.get("username", username) # Get username from payload

                if msg_type == "chat_message":
                    chat_text = payload.get('message', ''); voice = payload.get('voice', FUN_USERNAMES.get(sender_username, "en-US-AriaNeural"));
                    run_async(save_chat_entry, sender_username, chat_text, voice) # Fire-and-forget
                    await broadcast_message(message, exclude_id=client_id) # Forward

                elif msg_type == "place_object":
                    obj_data = payload.get("object_data");
                    if obj_data and 'obj_id' in obj_data and 'type' in obj_data:
                        with world_objects_lock: world_objects[obj_data['obj_id']] = obj_data
                        broadcast_payload = json.dumps({"type": "object_placed", "payload": {"object_data": obj_data, "username": sender_username}});
                        await broadcast_message(broadcast_payload, exclude_id=client_id)
                    else: print(f"WS Invalid place_object payload: {payload}")

                elif msg_type == "delete_object":
                     obj_id = payload.get("obj_id"); removed = False
                     if obj_id:
                          with world_objects_lock:
                              if obj_id in world_objects: del world_objects[obj_id]; removed = True
                          if removed:
                               broadcast_payload = json.dumps({"type": "object_deleted", "payload": {"obj_id": obj_id, "username": sender_username}});
                               await broadcast_message(broadcast_payload, exclude_id=client_id)
                     else: print(f"WS Invalid delete_object payload: {payload}")

                elif msg_type == "player_position":
                     pos_data = payload.get("position"); rot_data = payload.get("rotation")
                     if pos_data:
                          broadcast_payload = json.dumps({"type": "player_moved", "payload": {"username": sender_username, "id": client_id, "position": pos_data, "rotation": rot_data}});
                          await broadcast_message(broadcast_payload, exclude_id=client_id)

            except json.JSONDecodeError: print(f"WS Invalid JSON from {client_id}: {message[:100]}...")
            except Exception as e: print(f"WS Error processing msg from {client_id}: {e}")
    except websockets.ConnectionClosed: print(f"WS Client disconnected: {client_id} ({username})")
    except Exception as e: print(f"WS Unexpected handler error {client_id}: {e}")
    finally:
         await broadcast_message(json.dumps({"type": "user_leave", "payload": {"username": username, "id": client_id}}), exclude_id=client_id);
         await unregister_client(websocket) # Cleanup


async def run_websocket_server():
    """Coroutine to run the WebSocket server."""
    if st.session_state.get('server_running_flag', False): return
    st.session_state['server_running_flag'] = True; print("Attempting start WS server 0.0.0.0:8765...")
    stop_event = asyncio.Event(); st.session_state['websocket_stop_event'] = stop_event
    server = None
    try:
        # Changed host to 0.0.0.0 for accessibility, ensure firewall allows port 8765
        server = await websockets.serve(websocket_handler, "0.0.0.0", 8765); st.session_state['server_instance'] = server
        print(f"WS server started: {server.sockets[0].getsockname()}. Waiting for stop signal...")
        await stop_event.wait() # Keep running
    except OSError as e: print(f"### FAILED START WS SERVER: {e}"); st.session_state['server_running_flag'] = False; # Reset flag on failure
    except Exception as e: print(f"### UNEXPECTED WS SERVER ERROR: {e}"); st.session_state['server_running_flag'] = False; # Reset flag on failure
    finally:
        print("WS server task finishing...");
        if server: server.close(); await server.wait_closed(); print("WS server closed.")
        st.session_state['server_running_flag'] = False; st.session_state['server_instance'] = None; st.session_state['websocket_stop_event'] = None

def start_websocket_server_thread():
    """Starts the WebSocket server in a separate thread if not already running."""
    if st.session_state.get('server_task') and st.session_state.server_task.is_alive(): return
    if st.session_state.get('server_running_flag', False): return
    print("Creating/starting new server thread.");
    def run_loop():
        current_loop = None
        try:
            current_loop = asyncio.get_event_loop()
            if current_loop.is_running():
                 print("Server thread: Attaching to existing running loop (rare case).")
                 # If already running, might need different approach, but usually new thread = new loop
                 # This case is less likely with daemon threads starting fresh.
            else:
                 raise RuntimeError("No running loop found initially - expected.")
        except RuntimeError: # No loop in this thread, create new one
             print("Server thread: Creating new event loop.")
             loop = asyncio.new_event_loop(); asyncio.set_event_loop(loop)
             try: loop.run_until_complete(run_websocket_server())
             finally:
                  # Gracefully shutdown tasks if loop is closing
                  tasks = asyncio.all_tasks(loop)
                  for task in tasks: task.cancel()
                  loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=True))
                  loop.close(); print("Server thread loop closed.")

    st.session_state.server_task = threading.Thread(target=run_loop, daemon=True); st.session_state.server_task.start(); time.sleep(1.5)
    if not st.session_state.server_task.is_alive(): print("### Server thread failed to stay alive!")


# ==============================================================================
# Streamlit UI Layout Functions
# ==============================================================================

def render_sidebar():
    """Renders the Streamlit sidebar contents."""
    with st.sidebar:
        st.header("💾 World Versions")
        st.caption("Load or save named world states.")

        saved_worlds = get_saved_worlds()
        # Format options for radio button display
        world_options_display = {os.path.basename(w['filename']): f"{w['name']} ({w['timestamp']})" for w in saved_worlds}
        # The actual options list stores basenames
        radio_options_basenames = [None] + [os.path.basename(w['filename']) for w in saved_worlds]

        current_selection_basename = st.session_state.get('current_world_file', None)
        current_radio_index = 0
        if current_selection_basename and current_selection_basename in radio_options_basenames:
             try: current_radio_index = radio_options_basenames.index(current_selection_basename)
             except ValueError: current_radio_index = 0 # Default to None if not found

        selected_basename = st.radio(
             "Load World:", options=radio_options_basenames, index=current_radio_index,
             format_func=lambda x: "Live State (Unsaved)" if x is None else world_options_display.get(x, x), # Display formatted name
             key="world_selector_radio"
         )

        # Handle selection change
        if selected_basename != current_selection_basename:
             st.session_state.current_world_file = selected_basename # Store selected basename
             if selected_basename:
                  with st.spinner(f"Loading {selected_basename}..."):
                      if load_world_state_from_md(selected_basename):
                          run_async(broadcast_world_update) # Broadcast new state
                          st.toast("World loaded!", icon="✅")
                      else: st.error("Failed to load world."); st.session_state.current_world_file = None
             else: # Switched to "Live State"
                  print("Switched to live state.")
                  # Optionally clear world state or just stop tracking file? Stop tracking.
                  # Maybe broadcast current live state to ensure consistency?
                  # run_async(broadcast_world_update)
                  st.toast("Switched to Live State.")
             st.rerun()

        # Download Links for Worlds
        st.caption("Download:")
        cols = st.columns([4, 1]) # Columns for name and download button
        with cols[0]: st.write("**World Name** (Timestamp)")
        with cols[1]: st.write("**DL**")
        # Display max 10 worlds initially, add expander if more?
        for world_info in saved_worlds[:10]: # Limit display
            f_basename = os.path.basename(world_info['filename'])
            f_fullpath = os.path.join(SAVED_WORLDS_DIR, f_basename) # Reconstruct full path for link
            display_name = world_info.get('name', f_basename)
            timestamp = world_info.get('timestamp', 'N/A')
            col1, col2 = st.columns([4, 1])
            with col1: st.write(f"<small>{display_name} ({timestamp})</small>", unsafe_allow_html=True)
            with col2: st.markdown(get_download_link(f_fullpath, "md"), unsafe_allow_html=True)
        if len(saved_worlds) > 10:
             with st.expander(f"Show {len(saved_worlds)-10} more..."):
                  for world_info in saved_worlds[10:]:
                      # Repeat display logic
                      f_basename = os.path.basename(world_info['filename'])
                      f_fullpath = os.path.join(SAVED_WORLDS_DIR, f_basename)
                      display_name = world_info.get('name', f_basename)
                      timestamp = world_info.get('timestamp', 'N/A')
                      col1, col2 = st.columns([4, 1])
                      with col1: st.write(f"<small>{display_name} ({timestamp})</small>", unsafe_allow_html=True)
                      with col2: st.markdown(get_download_link(f_fullpath, "md"), unsafe_allow_html=True)

        st.markdown("---")

        # Build Tools Section
        st.header("🏗️ Build Tools")
        st.caption("Select an object to place.")
        cols = st.columns(5)
        col_idx = 0
        current_tool = st.session_state.get('selected_object', 'None')
        for emoji, name in PRIMITIVE_MAP.items():
            button_key = f"primitive_{name}"; button_type = "primary" if current_tool == name else "secondary"
            if cols[col_idx % 5].button(emoji, key=button_key, help=name, type=button_type, use_container_width=True):
                if st.session_state.get('selected_object', 'None') != name:
                    st.session_state.selected_object = name
                    # Fire and forget JS update
                    run_async(lambda name_arg=name: streamlit_js_eval(f"updateSelectedObjectType({json.dumps(name_arg)});", key=f"update_tool_js_{name_arg}"))
                    st.rerun()
            col_idx += 1
        st.markdown("---")
        if st.button("🚫 Clear Tool", key="clear_tool", use_container_width=True):
             if st.session_state.get('selected_object', 'None') != 'None':
                  st.session_state.selected_object = 'None'
                  run_async(lambda: streamlit_js_eval("updateSelectedObjectType('None');", key="update_tool_js_none"))
                  st.rerun()

        # Voice/User Section
        st.markdown("---")
        st.header("🗣️ Voice & User")
        current_username = st.session_state.get('username', list(FUN_USERNAMES.keys())[0])
        username_options = list(FUN_USERNAMES.keys()); current_index = 0
        try: current_index = username_options.index(current_username)
        except ValueError: current_index = 0 # Handle case where saved username is no longer valid
        new_username = st.selectbox("Change Name/Voice", options=username_options, index=current_index, key="username_select", format_func=lambda x: x.split(" ")[0])
        if new_username != st.session_state.username:
            old_username = st.session_state.username
            change_msg = json.dumps({"type":"user_rename", "payload": {"old_username": old_username, "new_username": new_username}})
            run_async(broadcast_message, change_msg) # Fire and forget broadcast
            st.session_state.username = new_username; st.session_state.tts_voice = FUN_USERNAMES[new_username]; save_username(st.session_state.username)
            st.rerun()
        st.session_state['enable_audio'] = st.toggle("Enable TTS Audio", value=st.session_state.get('enable_audio', True))


def render_main_content():
    """Renders the main content area with tabs."""
    st.title(f"{Site_Name} - User: {st.session_state.username}")

    tab_world, tab_chat, tab_pdf, tab_files = st.tabs(["🏗️ World Builder", "🗣️ Chat", "📚 PDF Tools", "📂 Files & Settings"])

    # --- World Builder Tab ---
    with tab_world:
        st.header("Shared 3D World")
        st.caption("Place objects using the sidebar tools. Changes are shared live!")
        current_file_basename = st.session_state.get('current_world_file', None)
        if current_file_basename:
             # Reconstruct full path for parsing if needed, or just use basename
             parsed = parse_world_filename(os.path.join(SAVED_WORLDS_DIR, current_file_basename)) # Parse info
             st.info(f"Current World: **{parsed['name']}** (`{current_file_basename}`)")
        else: st.info("Live State Active (Unsaved changes only persist if saved)")

        # Embed HTML Component
        html_file_path = 'index.html'
        try:
            with open(html_file_path, 'r', encoding='utf-8') as f: html_template = f.read()
            ws_url = "ws://localhost:8765" # Default for local dev
            try: # Attempt to get dynamic host
                 from streamlit.web.server.server import Server
                 session_info = Server.get_current()._get_session_info(st.runtime.scriptrunner.get_script_run_ctx().session_id)
                 server_host = session_info.ws.stream.request.host.split(':')[0]
                 ws_url = f"ws://{server_host}:8765"
            except Exception as e: print(f"WS URL detection failed ({e}), using localhost.")

            js_injection_script = f"""<script>
    window.USERNAME = {json.dumps(st.session_state.username)};
    window.WEBSOCKET_URL = {json.dumps(ws_url)};
    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:", {{ username: window.USERNAME, websocketUrl: window.WEBSOCKET_URL, selectedObject: window.SELECTED_OBJECT_TYPE }});
</script>"""
            html_content_with_state = html_template.replace('</head>', js_injection_script + '\n</head>', 1)
            components.html(html_content_with_state, height=700, scrolling=False)
        except FileNotFoundError: st.error(f"CRITICAL ERROR: Could not find '{html_file_path}'.")
        except Exception as e: st.error(f"Error loading 3D component: {e}"); st.exception(e)

    # --- Chat Tab ---
    with tab_chat:
        st.header(f"{START_ROOM} Chat")
        # Load history - use run_async result or session state if already loaded
        if 'chat_history' not in st.session_state or not st.session_state.chat_history:
             chat_history = asyncio.run(load_chat_history()) # Blocking load if first time
        else:
             chat_history = st.session_state.chat_history

        chat_container = st.container(height=500)
        with chat_container:
             if chat_history: st.markdown("----\n".join(reversed(chat_history[-50:]))) # Show last 50, reversed
             else: st.caption("No chat messages yet.")

        # Chat Input Area
        message_value = st.text_input("Your Message:", key="message_input", label_visibility="collapsed")
        send_button_clicked = st.button("Send Chat 💬", key="send_chat_button")
        should_autosend = st.session_state.get('autosend', False) and message_value

        if send_button_clicked or should_autosend:
            message_to_send = message_value
            if message_to_send.strip() and message_to_send != st.session_state.get('last_message', ''):
                st.session_state.last_message = message_to_send # Update tracker
                voice = FUN_USERNAMES.get(st.session_state.username, "en-US-AriaNeural")
                ws_message = json.dumps({"type": "chat_message", "payload": {"username": st.session_state.username, "message": message_to_send, "voice": voice}})
                # Fire and forget async tasks
                run_async(broadcast_message, ws_message)
                run_async(save_chat_entry, st.session_state.username, message_to_send, voice)
                st.session_state.message_input = "" # Clear state for next run
                st.rerun()
            elif send_button_clicked: st.toast("Message empty or same as last.")
        st.checkbox("Autosend Chat", key="autosend") # Toggle autosend

    # --- PDF Tab ---
    with tab_pdf:
         st.header("📚 PDF Tools")
         pdf_file = st.file_uploader("Upload PDF for Audio Conversion", type="pdf", key="pdf_upload")
         max_pages = st.slider('Max Pages to Process', 1, 50, 10, key="pdf_pages")
         if pdf_file:
              # Use a button to trigger potentially long processing
              if st.button("Process PDF to Audio", key="process_pdf_button"):
                   with st.spinner("Processing PDF... This may take time."):
                        process_pdf_tab(pdf_file, max_pages, st.session_state.tts_voice)

    # --- Files & Settings Tab ---
    with tab_files:
        st.header("📂 Files & Settings")
        st.subheader("💾 World State Management")
        current_file_basename = st.session_state.get('current_world_file', None)

        # Save Current Version Button
        if current_file_basename:
             parsed = parse_world_filename(os.path.join(SAVED_WORLDS_DIR, current_file_basename))
             save_label = f"Save Changes to '{parsed['name']}'"
             if st.button(save_label, key="save_current_world", help=f"Overwrite '{current_file_basename}'"):
                 with st.spinner(f"Overwriting {current_file_basename}..."):
                     if save_world_state_to_md(current_file_basename): st.success("Current world saved!")
                     else: st.error("Failed to save world.")
         else:
             st.info("Load a world from the sidebar to enable saving changes to it.")

        # Save As New Version Section
        st.subheader("Save As New Version")
        new_name_files = st.text_input("New World Name:", key="new_world_name_files", value=st.session_state.get('new_world_name', 'MyWorld'))
        if st.button("💾 Save Live State as New Version", key="save_new_version_files"):
             if new_name_files.strip():
                  new_filename_base = generate_world_save_filename(new_name_files)
                  with st.spinner(f"Saving new version '{new_name_files}'..."):
                      if save_world_state_to_md(new_filename_base):
                          st.success(f"Saved as {new_filename_base}")
                          st.session_state.current_world_file = new_filename_base # Switch to new file
                          st.session_state.new_world_name = "MyWorld" # Reset default for next time
                          st.rerun()
                      else: st.error("Failed to save new version.")
             else: st.warning("Please enter a name.")

        st.subheader("⚙️ Server Status")
        col_ws, col_clients = st.columns(2)
        with col_ws:
             server_alive = st.session_state.get('server_task') and st.session_state.server_task.is_alive(); ws_status = "Running" if server_alive else "Stopped"; st.metric("WebSocket Server", ws_status)
             if not server_alive and st.button("Restart Server Thread", key="restart_ws"): start_websocket_server_thread(); st.rerun()
        with col_clients: st.metric("Connected Clients", len(connected_clients))

        # File Deletion
        st.subheader("🗑️ Delete Files")
        st.warning("Deletion is permanent!", icon="⚠️")
        col_del1, col_del2, col_del3, col_del4 = st.columns(4)
        with col_del1:
             if st.button("🗑️ Chats", key="del_chat_md"): delete_files([os.path.join(CHAT_DIR, "*.md")]); st.session_state.chat_history = []; st.rerun()
        with col_del2:
            if st.button("🗑️ Audio", key="del_audio_mp3"): delete_files([os.path.join(AUDIO_DIR, "*.mp3"), os.path.join(AUDIO_CACHE_DIR, "*.mp3")]); st.session_state.audio_cache = {}; st.rerun()
        with col_del3:
            if st.button("🗑️ Worlds", key="del_worlds_md"): delete_files([os.path.join(SAVED_WORLDS_DIR, f"{WORLD_STATE_FILE_MD_PREFIX}*.md")], exclude_files=[st.session_state.get('current_world_file')]); st.session_state.current_world_file = None; st.rerun() # Protect current? No, delete all.
        with col_del4:
            if st.button("🗑️ All Gen", key="del_all_gen", help="Deletes Chats, Audio, Worlds, Zips"): delete_files([os.path.join(CHAT_DIR, "*.md"), os.path.join(AUDIO_DIR, "*.mp3"), os.path.join(AUDIO_CACHE_DIR, "*.mp3"), os.path.join(SAVED_WORLDS_DIR, "*.md"), "*.zip"]); st.session_state.chat_history = []; st.session_state.audio_cache = {}; st.session_state.current_world_file = None; st.rerun()

        # Download Archives
        st.subheader("📦 Download Archives")
        zip_files = sorted(glob.glob(os.path.join(MEDIA_DIR,"*.zip")), key=os.path.getmtime, reverse=True) # Look in base dir
        if zip_files:
             # Zip specific content types
             col_zip1, col_zip2, col_zip3 = st.columns(3)
             with col_zip1:
                  if st.button("Zip Worlds"): create_zip_of_files(glob.glob(os.path.join(SAVED_WORLDS_DIR, "*.md")), "Worlds")
             with col_zip2:
                   if st.button("Zip Chats"): create_zip_of_files(glob.glob(os.path.join(CHAT_DIR, "*.md")), "Chats")
             with col_zip3:
                   if st.button("Zip Audio"): create_zip_of_files(glob.glob(os.path.join(AUDIO_DIR, "*.mp3")) + glob.glob(os.path.join(AUDIO_CACHE_DIR, "*.mp3")), "Audio")

             st.caption("Existing Zip Files:")
             for zip_file in zip_files: st.markdown(get_download_link(zip_file, "zip"), unsafe_allow_html=True)
        else: st.caption("No zip archives found.")


# ==============================================================================
# Main Execution Logic
# ==============================================================================

def initialize_world():
    """Loads initial world state (most recent) if not already loaded."""
    # This check prevents reloading state on every single rerun, only on session start
    if not st.session_state.get('initial_world_state_loaded', False):
        print("Performing initial world load for session...")
        saved_worlds = get_saved_worlds()
        loaded_successfully = False
        if saved_worlds:
             # Load the most recent world file (first in sorted list)
             latest_world_file_basename = os.path.basename(saved_worlds[0]['filename'])
             print(f"Loading most recent world on startup: {latest_world_file_basename}")
             if load_world_state_from_md(latest_world_file_basename): # This updates global state and sets session state 'current_world_file'
                  loaded_successfully = True
             else:
                  print("Failed to load most recent world, starting empty.")
         else:
              print("No saved worlds found, starting with empty state.")

        # Ensure global dict is empty if no file loaded successfully
        if not loaded_successfully:
             with world_objects_lock: world_objects.clear()
             st.session_state.current_world_file = None # Ensure no file is marked as loaded

        st.session_state.initial_world_state_loaded = True # Mark as loaded for this session
        print("Initial world load process complete.")

if __name__ == "__main__":
    # 1. Initialize session state first (essential for other checks)
    init_session_state()

    # 2. Start WebSocket server thread if needed (check flags and thread life)
    # Use server_running_flag to prevent multiple start attempts
    server_thread = st.session_state.get('server_task')
    server_alive = server_thread is not None and server_thread.is_alive()
    if not st.session_state.get('server_running_flag', False) and not server_alive:
         start_websocket_server_thread()
    elif server_alive and not st.session_state.get('server_running_flag', False):
         # Correct flag if thread is alive but flag is false
         st.session_state.server_running_flag = True

    # 3. Load initial world state from disk if not already done for this session
    initialize_world()

    # 4. Render the UI (Sidebar and Main Content)
    render_sidebar()
    render_main_content()

    # 5. Optional Periodic Save (Example - uncomment to enable)
    # interval_seconds = 300 # 5 minutes
    # if 'last_periodic_save' not in st.session_state: st.session_state.last_periodic_save = 0
    # if time.time() - st.session_state.last_periodic_save > interval_seconds:
    #     current_file_to_save = st.session_state.get('current_world_file')
    #     if current_file_to_save: # Only save if a specific file is loaded
    #         print(f"Triggering periodic save for {current_file_to_save}...")
    #         if save_world_state_to_md(current_file_to_save):
    #             st.session_state.last_periodic_save = time.time()
    #             print("Periodic save successful.")
    #         else:
    #             print("Periodic save failed.")
    #     else:
    #          st.session_state.last_periodic_save = time.time() # Reset timer even if not saving