# app.py (Re-integrated WebSockets for 3D Sync - Cleaned)
import streamlit as st
import asyncio
import websockets # Re-added
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, deque
from streamlit_js_eval import streamlit_js_eval # Keep for UI interaction if needed
from PIL import Image

# ==============================================================================
# 1. ⚙️ Configuration & Constants
# ==============================================================================

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

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

# General Constants
Site_Name = '🏗️ Live World Builder ⚡'
MEDIA_DIR = "."
STATE_FILE = "user_state.txt"
DEFAULT_TTS_VOICE = "en-US-AriaNeural"

# Directories
CHAT_DIR = "chat_logs"
AUDIO_CACHE_DIR = "audio_cache"
AUDIO_DIR = "audio_logs"
SAVED_WORLDS_DIR = "saved_worlds"

# World Builder Constants
PLOT_WIDTH = 50.0
PLOT_DEPTH = 50.0
WORLD_STATE_FILE_MD_PREFIX = "🌌_" # Keep prefix for saved files
MAX_ACTION_LOG_SIZE = 30

# 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",
} # Simplified list
EDGE_TTS_VOICES = list(set(FUN_USERNAMES.values()))

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

# Primitives Map
PRIMITIVE_MAP = {
    "Tree": "🌳", "Rock": "🗿", "Simple House": "🏛️", "Pine Tree": "🌲", "Brick Wall": "🧱",
    "Sphere": "🔵", "Cube": "📦", "Cylinder": "🧴", "Cone": "🍦", "Torus": "🍩",
    "Mushroom": "🍄", "Cactus": "🌵", "Campfire": "🔥", "Star": "⭐", "Gem": "💎",
    "Tower": "🗼", "Barrier": "🚧", "Fountain": "⛲", "Lantern": "🏮", "Sign Post": "팻"
}
TOOLS_MAP = {"None": "🚫"}
TOOLS_MAP.update({name: emoji for emoji, name in PRIMITIVE_MAP.items()})

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

# --- Global State (WebSocket Client Tracking Only) ---
clients_lock = threading.Lock()
connected_clients = set() # Holds client_id strings (websocket.id)

# ==============================================================================
# 2. ✨ 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=25):
    """Cleans a string part for use in a filename."""
    if not isinstance(text, str): text = "invalid_name"
    text = re.sub(r'\s+', '_', text)
    text = re.sub(r'[^\w\-.]', '', text)
    return text[:max_len]

def run_async(async_func, *args, **kwargs):
    """Runs an async function safely from a sync context using create_task or asyncio.run."""
    try:
        loop = asyncio.get_running_loop()
        return loop.create_task(async_func(*args, **kwargs))
    except RuntimeError:
        try: return asyncio.run(async_func(*args, **kwargs))
        except Exception as e: print(f"❌ Error run_async new loop: {e}"); return None
    except Exception as e: print(f"❌ Error run_async schedule task: {e}"); return None

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

# ==============================================================================
# 3. 🌍 World State Manager (Using st.cache_resource)
# ==============================================================================

def get_saved_worlds(): # Define this before it's used in load_initial_world_from_file
    """Scans the saved worlds directory for world MD files and parses them."""
    try:
        ensure_dir(SAVED_WORLDS_DIR);
        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] # parse_world_filename needs to be defined below
        parsed_worlds.sort(key=lambda x: x.get('dt') if x.get('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 []

def parse_world_filename(filename): # Define this before get_saved_worlds uses it indirectly via load_initial
    """Extracts info from filename if possible, otherwise returns defaults."""
    basename = os.path.basename(filename)
    if basename.startswith(WORLD_STATE_FILE_MD_PREFIX) and basename.endswith(".md"):
        core = basename[len(WORLD_STATE_FILE_MD_PREFIX):-3]; parts = core.split('_')
        if len(parts) >= 5 and parts[-3] == "by":
            timestamp_str = parts[-2]; username = parts[-4]; world_name = " ".join(parts[:-4]); dt_obj = None
            try: dt_obj = pytz.utc.localize(datetime.strptime(timestamp_str, '%Y%m%d_%H%M%S'))
            except Exception: dt_obj = None
            return {"name": world_name or "Untitled", "user": username, "timestamp": timestamp_str, "dt": dt_obj, "filename": filename}
    # Fallback
    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','').replace(WORLD_STATE_FILE_MD_PREFIX, ''), "user": "Unknown", "timestamp": "Unknown", "dt": dt_fallback, "filename": filename}

def load_initial_world_from_file():
    """Loads the state from the most recent MD file found."""
    print(f"[{time.time():.1f}] ⏳ Attempting to load initial world state from files...")
    loaded_state = defaultdict(dict)
    saved_worlds = get_saved_worlds()
    if saved_worlds:
        latest_world_file_basename = os.path.basename(saved_worlds[0]['filename'])
        print(f"⏳ Found most recent file: {latest_world_file_basename}")
        load_path = os.path.join(SAVED_WORLDS_DIR, latest_world_file_basename)
        if os.path.exists(load_path):
            try:
                with open(load_path, 'r', encoding='utf-8') as f: content = f.read()
                json_match = re.search(r"```json\s*(\{[\s\S]*?\})\s*```", content, re.IGNORECASE)
                if json_match:
                    world_data_dict = json.loads(json_match.group(1))
                    for k, v in world_data_dict.items(): loaded_state[str(k)] = v
                    print(f"✅ Successfully loaded {len(loaded_state)} objects for initial state.")
                    # Store the initially loaded file basename in session state here?
                    st.session_state._initial_world_file_loaded = latest_world_file_basename
                else: print("⚠️ No JSON block found in initial file.")
            except Exception as e: print(f"❌ Error parsing initial world file {latest_world_file_basename}: {e}")
        else: print(f"⚠️ Most recent file {latest_world_file_basename} not found at path {load_path}.")
    else: print("🌫️ No saved world files found to load initial state.")
    return loaded_state

@st.cache_resource(ttl=3600) # Cache resource for 1 hour
def get_world_state_manager():
    """
    Initializes and returns the shared world state dictionary and its lock.
    Loads initial state from the most recent file on first creation.
    """
    print(f"[{time.time():.1f}] --- ✨ Initializing/Retrieving Shared World State Resource ---")
    manager = {
        "lock": threading.Lock(),
        "state": load_initial_world_from_file() # Load initial state here
    }
    # Initial current_world_file is now handled after init_session_state in main logic
    return manager

def get_current_world_state_copy():
    """Safely gets a copy of the current world state dictionary."""
    manager = get_world_state_manager()
    with manager["lock"]:
        return dict(manager["state"]) # Return a copy

# ==============================================================================
# 4. 💾 World State File Handling (Save/Load - Refactored for Cached State)
# ==============================================================================
def generate_world_save_filename(username="User", world_name="World"):
    timestamp = get_current_time_str(); clean_user = clean_filename_part(username, 15);
    clean_world = clean_filename_part(world_name, 20);
    rand_hash = hashlib.md5(str(time.time()).encode()+username.encode()+world_name.encode()).hexdigest()[:4]
    return f"{WORLD_STATE_FILE_MD_PREFIX}{clean_world}_by_{clean_user}_{timestamp}_{rand_hash}.md"

def save_world_state_to_md(target_filename_base):
    """Saves the current cached world state to a specific MD file."""
    manager = get_world_state_manager()
    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 manager["lock"]:
        world_data_dict = dict(manager["state"])
        print(f"💾 Saving {len(world_data_dict)} objects...")
        parsed_info = parse_world_filename(save_path)
        timestamp_save = get_current_time_str()
        md_content = f"""# World State: {parsed_info['name']} by {parsed_info['user']}
* **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}")
    return success

def load_world_state_from_md(filename_base):
    """Loads world state from MD, updates cached state, returns success bool."""
    manager = get_world_state_manager()
    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()
        json_match = re.search(r"```json\s*(\{[\s\S]*?\})\s*```", content, re.IGNORECASE)
        if not json_match: st.error(f"Could not find JSON block in {filename_base}"); return False
        world_data_dict = json.loads(json_match.group(1))

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

    except json.JSONDecodeError as e: st.error(f"Invalid JSON 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

# ==============================================================================
# 5. 👤 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), # Stores websocket objects by client_id
        'last_chat_update': 0, 'message_input': "", 'audio_cache': {},
        'tts_voice': DEFAULT_TTS_VOICE, 'chat_history': [], 'enable_audio': True,
        'download_link_cache': {}, 'username': None, 'autosend': False,
        'last_message': "",
        'selected_object': 'None',
        # 'initial_world_state_loaded' flag removed, cache resource handles init
        'current_world_file': None, # Track loaded world filename (basename)
        'new_world_name': "MyDreamscape",
        'action_log': deque(maxlen=MAX_ACTION_LOG_SIZE),
        # State related to JS interaction moved or removed if WS handles it
    }
    for k, v in defaults.items():
        if k not in st.session_state:
            if isinstance(v, (deque, dict, list)): st.session_state[k] = v.copy()
            else: st.session_state[k] = v
    # Ensure complex types are correctly initialized
    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.audio_cache, dict): st.session_state.audio_cache = {}
    if not isinstance(st.session_state.download_link_cache, dict): st.session_state.download_link_cache = {}
    if not isinstance(st.session_state.action_log, deque): st.session_state.action_log = deque(maxlen=MAX_ACTION_LOG_SIZE)

# ==============================================================================
# 6. 📝 Action Log Helper
# ==============================================================================
def add_action_log(message, emoji="➡️"):
    """Adds a timestamped message with emoji to the session's action log."""
    if 'action_log' not in st.session_state or not isinstance(st.session_state.action_log, deque):
        st.session_state.action_log = deque(maxlen=MAX_ACTION_LOG_SIZE)
    timestamp = datetime.now().strftime("%H:%M:%S")
    st.session_state.action_log.appendleft(f"{emoji} [{timestamp}] {message}")

# ==============================================================================
# 7. 🎧 Audio / TTS / Chat / File Handling Helpers
# ==============================================================================
# (Keep implementations from previous correct version - Placeholder for brevity)
def clean_text_for_tts(text): # ... implementation ...
    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): # ... implementation ...
    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); 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"): # ... implementation ...
    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)
            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>")
async def async_edge_tts_generate(text, voice, username): # ... implementation ...
    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): # ... implementation ...
    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}")
async def save_chat_entry(username, message, voice, is_markdown=False): # ... implementation ...
    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)
    if 'chat_history' not in st.session_state: st.session_state.chat_history = [];
    st.session_state.chat_history.append(entry)
    audio_file = None;
    if st.session_state.get('enable_audio', True): tts_message = message ; audio_file = await async_edge_tts_generate(tts_message, voice, username)
    return md_file, audio_file
async def load_chat_history(): # ... implementation ...
     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
          print(f"✅ Loaded {loaded_count} chat entries from files.")
     return st.session_state.chat_history
def create_zip_of_files(files_to_zip, prefix="Archive"): # ... implementation ...
    if not files_to_zip: st.warning("💨 Nothing to gather into an archive."); return None
    timestamp = format_timestamp_prefix(f"Zip_{prefix}"); zip_name = f"{prefix}_{timestamp}.zip"
    try:
        print(f"📦 Creating zip archive: {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 file: {f}")
        print("✅ Zip archive created successfully."); st.success(f"Created {zip_name}"); return zip_name
    except Exception as e: print(f"❌ Zip creation failed: {e}"); st.error(f"Zip creation failed: {e}"); return None
def delete_files(file_patterns, exclude_files=None): # ... implementation ...
    protected = [STATE_FILE, "app.py", "index.html", "requirements.txt", "README.md"]
    current_world_base = st.session_state.get('current_world_file')
    if current_world_base: protected.append(current_world_base)
    if exclude_files: protected.extend(exclude_files)
    deleted_count = 0; errors = 0
    for pattern in file_patterns:
         pattern_path = pattern
         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)
                   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
                   #else: print(f"🚫 Skipping protected/directory: {f_path}") # Debugging
         except Exception as glob_e: print(f"❌ Error matching pattern {pattern}: {glob_e}"); errors += 1
    msg = f"✅ Successfully deleted {deleted_count} files." if errors == 0 and deleted_count > 0 else 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 unprotected files found to delete.")
    st.session_state['download_link_cache'] = {}; st.session_state['audio_cache'] = {}
async def save_pasted_image(image, username): # ... implementation ...
     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(): # ... implementation ...
      pasted_img = None; img_type = None
      paste_input_value = st.text_area("📋 Paste Image Data Here", key="paste_input_area", height=50, value=st.session_state.get('paste_image_base64_input', ""), help="Paste image data directly (e.g., from clipboard)")
      if st.button("🖼️ Process Pasted Image", key="process_paste_button"):
           st.session_state.paste_image_base64_input = paste_input_value
           if paste_input_value and paste_input_value.startswith('data:image'):
               try:
                   mime_type = paste_input_value.split(';')[0].split(':')[1]; base64_str = paste_input_value.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
                   st.session_state.paste_image_base64_input = ""
                   st.rerun()
               except ImportError: st.error("⚠️ Pillow library needed.")
               except Exception as e: st.error(f"❌ Img decode err: {e}"); st.session_state.paste_image_base64 = ""; st.session_state.paste_image_base64_input = paste_input_value
           else: st.warning("⚠️ No valid image data pasted."); st.session_state.paste_image_base64 = ""; st.session_state.paste_image_base64_input = paste_input_value
      processed_b64 = st.session_state.get('paste_image_base64', '')
      if processed_b64:
           try: img_bytes = base64.b64decode(processed_b64); return Image.open(io.BytesIO(img_bytes))
           except Exception: return None
      return None
class AudioProcessor: # ... implementation ...
      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=DEFAULT_TTS_VOICE):
            cache_key=hashlib.md5(f"{text[:150]}:{voice}".encode()).hexdigest(); cache_path=os.path.join(self.cache_dir, f"{cache_key}.mp3");
            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): # ... implementation ...
    st.subheader("📜 PDF Processing Results")
    if pdf_file is None: st.info("⬆️ Upload a PDF file and click 'Process PDF' to begin."); return
    audio_processor = AudioProcessor()
    try:
        reader=PdfReader(pdf_file);
        if reader.is_encrypted: st.warning("🔒 PDF is encrypted."); return
        total_pages_in_pdf = len(reader.pages); pages_to_process = min(total_pages_in_pdf, max_pages);
        st.write(f"⏳ Processing first {pages_to_process} of {total_pages_in_pdf} pages from '{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:
                 audio_path = asyncio.run(run_async_audio()) # Use asyncio.run in 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}")

        for i in range(pages_to_process):
             try: # Start try block for page processing
                 page = reader.pages[i]
                 text = page.extract_text() # Attempt text extraction
                 if text and text.strip(): # Check extracted text
                     texts[i]=text # Store text
                     thread = threading.Thread(target=process_page_sync, args=(i, text)) # Create thread
                     page_threads.append(thread) # Append thread
                     thread.start() # Start thread
                 else: # Handle empty extraction
                     texts[i] = "[📄 No text extracted or page empty]"
                     # print(f"Page {i+1}: No text extracted.") # Verbose
             # Correctly indented except block
             except Exception as extract_e:
                  texts[i] = f"[❌ Error extract: {extract_e}]" # Store error message
                  print(f"Error page {i+1} extract: {extract_e}") # Log error

        progress_bar = st.progress(0.0, text="✨ Transmuting pages to sound...")
        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), text=f"✨ Processed {completed_threads}/{total_threads} pages...")
            if time.time() - start_join_time > 600: print("⌛ PDF processing timed out."); st.warning("Processing timed out."); break
            time.sleep(0.5)
        progress_bar.progress(1.0, text="✅ Processing complete.")

        st.write("🎶 Results:")
        for i in range(pages_to_process):
              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:
                        page_text = texts.get(i,"")
                        if page_text.strip() and not page_text.startswith("["): st.caption("🔇 Audio generation failed or timed out.")
                        # else: st.caption("🔇 No text to generate audio from.") # Implicit

    except ImportError: st.error("⚠️ PyPDF2 library needed.")
    except Exception as pdf_e: st.error(f"❌ Error reading PDF '{pdf_file.name}': {pdf_e}"); st.exception(pdf_e)

# ==============================================================================
# 8. 🕸️ WebSocket Server Logic (Re-added for Chat/Presence)
# ==============================================================================

async def register_client(websocket):
    """Adds client to tracking structures, ensuring thread safety."""
    client_id = str(websocket.id);
    with clients_lock:
        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):
    """Removes client from tracking structures, ensuring thread safety."""
    client_id = str(websocket.id);
    with clients_lock:
        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 # Raise to be caught by gather
     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."""
    # Create local copies under lock for thread safety
    with clients_lock:
        if not connected_clients: return
        current_client_ids = list(connected_clients)
        # Ensure active_connections exists and make a copy
        if 'active_connections' in st.session_state:
            active_connections_copy = st.session_state.active_connections.copy()
        else:
            active_connections_copy = {} # Should not happen if init_session_state is correct

    tasks = []
    for client_id in current_client_ids:
        if client_id == exclude_id: continue
        websocket = active_connections_copy.get(client_id) # Use copy
        if websocket:
            tasks.append(asyncio.create_task(send_safely(websocket, message, client_id)))

    if tasks:
        results = await asyncio.gather(*tasks, return_exceptions=True)
        # Optional: Check results for exceptions if specific error handling per client is needed


async def broadcast_world_update():
    """Broadcasts the current world state (from cache) to all clients."""
    # Uses the cached state manager
    world_state_copy = get_current_world_state_copy()
    update_msg = json.dumps({"type": "initial_state", "payload": world_state_copy})
    print(f"📡 Broadcasting full world update ({len(world_state_copy)} objects)...")
    await broadcast_message(update_msg)

async def websocket_handler(websocket, path):
    """Handles WebSocket connections and messages (primarily for Chat & 3D Sync)."""
    await register_client(websocket); client_id = str(websocket.id);
    username = st.session_state.get('username', f"User_{client_id[:4]}")

    try: # Send initial world state
        initial_state_payload = get_current_world_state_copy() # Get state using cached helper
        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}")
        # Announce join after state sent
        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 during initial phase {client_id}: {e}")

    try: # Message processing 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

                # --- Handle Different Message Types ---
                manager = get_world_state_manager() # Get state manager for world updates

                if msg_type == "chat_message":
                    chat_text = payload.get('message', ''); voice = payload.get('voice', FUN_USERNAMES.get(sender_username, DEFAULT_TTS_VOICE));
                    print(f"💬 WS Recv Chat from {sender_username}: {chat_text[:30]}...")
                    run_async(save_chat_entry, sender_username, chat_text, voice) # Save locally async
                    await broadcast_message(message, exclude_id=client_id) # Broadcast chat

                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:
                        print(f"➕ WS Recv Place from {sender_username}: {obj_data['type']} ({obj_data['obj_id']})")
                        with manager["lock"]: manager["state"][obj_data['obj_id']] = obj_data # Update cached state
                        # Broadcast placement to others
                        broadcast_payload = json.dumps({"type": "object_placed", "payload": {"object_data": obj_data, "username": sender_username}});
                        await broadcast_message(broadcast_payload, exclude_id=client_id)
                        run_async(lambda: add_action_log(f"Placed {obj_data['type']} ({obj_data['obj_id'][:6]}) by {sender_username}", TOOLS_MAP.get(obj_data['type'], '❓')))
                    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:
                          print(f"➖ WS Recv Delete from {sender_username}: {obj_id}")
                          with manager["lock"]:
                              if obj_id in manager["state"]: del manager["state"][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)
                               run_async(lambda: add_action_log(f"Deleted obj ({obj_id[:6]}) by {sender_username}", "🗑️"))
                     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) # Broadcast movement

                elif msg_type == "ping": await websocket.send(json.dumps({"type": "pong"}))
                else: print(f"⚠️ WS Recv unknown type from {client_id}: {msg_type}")

            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)


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:
        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()
    except OSError as e: print(f"### ❌ FAILED START WS SERVER: {e}"); st.session_state['server_running_flag'] = False;
    except Exception as e: print(f"### ❌ UNEXPECTED WS SERVER ERROR: {e}"); st.session_state['server_running_flag'] = False;
    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(): # Wrapper to manage event loop in thread
        loop = None
        try: loop = asyncio.get_running_loop()
        except RuntimeError: loop = asyncio.new_event_loop(); asyncio.set_event_loop(loop)
        try: loop.run_until_complete(run_websocket_server())
        finally:
             if loop and not loop.is_closed():
                  tasks = asyncio.all_tasks(loop);
                  if tasks:
                    for task in tasks: task.cancel()
                    try: loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=True))
                    except asyncio.CancelledError: pass
                  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!")


# ==============================================================================
# 9. 🎨 Streamlit UI Layout Functions
# ==============================================================================

def render_sidebar():
    """Renders the Streamlit sidebar contents."""
    with st.sidebar:
        # 1. World Management
        st.header("1. 💾 World Management")
        st.caption("💾 Save the current view or ✨ load a past creation.")

        # World Save Button
        current_file = st.session_state.get('current_world_file')
        save_name_value = st.session_state.get('world_save_name_input', "MyDreamscape" if not current_file else parse_world_filename(os.path.join(SAVED_WORLDS_DIR, current_file)).get("name", current_file))
        world_save_name = st.text_input("World Name:", key="world_save_name_input", value=save_name_value, help="Enter name to save.")

        if st.button("💾 Save Current World View", key="sidebar_save_world"):
            if not world_save_name.strip(): st.warning("⚠️ Please enter a World Name.")
            else:
                # Save current state (which is managed by cache resource, updated by WS)
                filename_to_save = ""; is_overwrite = False
                if current_file:
                     try: # Check if name matches current loaded file's parsed name
                          parsed_current = parse_world_filename(os.path.join(SAVED_WORLDS_DIR, current_file))
                          if world_save_name == parsed_current.get('name', ''): filename_to_save = current_file; is_overwrite = True
                     except Exception: pass # Fallback to new save if parsing fails

                if not filename_to_save: filename_to_save = generate_world_save_filename(st.session_state.username, world_save_name)

                op_text = f"Overwriting {filename_to_save}..." if is_overwrite else f"Saving as {filename_to_save}..."
                with st.spinner(op_text):
                    if save_world_state_to_md(filename_to_save): # Saves state from cached resource
                        action = "Overwritten" if is_overwrite else "Saved new"
                        st.success(f"World {action}: {filename_to_save}"); add_action_log(f"Saved world: {filename_to_save}", emoji="💾")
                        st.session_state.current_world_file = filename_to_save # Track saved file
                        st.rerun()
                    else: st.error("❌ Failed to save world state.")


        # --- World Load ---
        st.markdown("---")
        st.header("2. 📂 Load World")
        st.caption("📜 Unfurl a previously woven dreamscape.")
        saved_worlds = get_saved_worlds()

        if not saved_worlds: st.caption("🌫️ The archives are empty.")
        else:
             cols_header = st.columns([4, 1, 1]);
             with cols_header[0]: st.write("**Name** (User, Time)")
             with cols_header[1]: st.write("**Load**")
             with cols_header[2]: st.write("**DL**")

             list_container = st.container(height=300 if len(saved_worlds) > 7 else None)
             with list_container:
                 for world_info in saved_worlds:
                    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); user = world_info.get('user', 'N/A'); timestamp = world_info.get('timestamp', 'N/A')
                    display_text = f"{display_name} ({user}, {timestamp})"
                    col1, col2, col3 = st.columns([4, 1, 1])
                    with col1: st.write(f"<small>{display_text}</small>", unsafe_allow_html=True)
                    with col2:
                        is_current = (st.session_state.get('current_world_file') == f_basename)
                        btn_load = st.button("✨", key=f"load_{f_basename}", help=f"Load {f_basename}", disabled=is_current)
                    with col3: st.markdown(get_download_link(f_fullpath, "md"), unsafe_allow_html=True)

                    if btn_load:
                        print(f"🖱️ Load button clicked for: {f_basename}")
                        with st.spinner(f"Loading {f_basename}..."):
                            # load_world_state_from_md now updates the cached resource directly
                            if load_world_state_from_md(f_basename):
                                run_async(broadcast_world_update) # Broadcast the newly loaded state
                                add_action_log(f"Loading world: {f_basename}", emoji="📂")
                                st.toast("World loaded!", icon="✅")
                                st.rerun() # Rerun to update UI and ensure clients get state via WS
                            else: st.error(f"❌ Failed to load world file: {f_basename}")

        # --- Build Tools ---
        st.markdown("---")
        st.header("3. 🛠️ Build Tools")
        st.caption("Select your creative instrument.")
        tool_options = list(TOOLS_MAP.keys())
        current_tool_name = st.session_state.get('selected_object', 'None')
        try: tool_index = tool_options.index(current_tool_name)
        except ValueError: tool_index = 0

        selected_tool = st.radio(
            "Select Tool:", options=tool_options, index=tool_index,
            format_func=lambda name: f"{TOOLS_MAP.get(name, '')} {name}",
            key="tool_selector_radio", horizontal=True, label_visibility="collapsed"
        )
        if selected_tool != current_tool_name:
            st.session_state.selected_object = selected_tool
            tool_emoji = TOOLS_MAP.get(selected_tool, '❓')
            add_action_log(f"Tool selected: {selected_tool}", emoji=tool_emoji)
            try: streamlit_js_eval(js_code=f"updateSelectedObjectType({json.dumps(selected_tool)});", key=f"update_tool_js_{selected_tool}")
            except Exception as e: print(f"❌ JS tool update error: {e}")
            st.rerun()


        # --- Action Log ---
        st.markdown("---")
        st.header("4. 📝 Action Log")
        st.caption("📜 A chronicle of your recent creative acts.")
        log_container = st.container(height=200)
        with log_container:
            log_entries = st.session_state.get('action_log', [])
            if log_entries: st.code('\n'.join(log_entries), language="log")
            else: st.caption("🌬️ The log awaits your first action...")


        # --- Voice/User ---
        st.markdown("---")
        st.header("5. 👤 Voice & User")
        st.caption("🎭 Choose your persona in this realm.")
        current_username = st.session_state.get('username', "DefaultUser")
        username_options = list(FUN_USERNAMES.keys()) if FUN_USERNAMES else [current_username]
        current_index = 0;
        try: # Safely find index
            if current_username in username_options: current_index = username_options.index(current_username)
        except ValueError: pass # Keep index 0

        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.get('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) # Broadcast name change
            st.session_state.username = new_username; st.session_state.tts_voice = FUN_USERNAMES.get(new_username, DEFAULT_TTS_VOICE); save_username(st.session_state.username)
            add_action_log(f"Persona changed to {new_username}", emoji="🎭")
            st.rerun()
        st.session_state['enable_audio'] = st.toggle("🔊 Enable TTS Audio", value=st.session_state.get('enable_audio', True), help="Generate audio for chat messages?")


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

    # NOTE: No longer need to check/send 'world_to_load_data' here.
    # The load button triggers load_world_state_from_md which updates the cache,
    # then triggers broadcast_world_update (via run_async), and reruns.
    # The WS handler sends initial state from the cache on new connections.

    # Define Tabs
    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 Dreamscape")
        st.caption("✨ Weave reality with sidebar tools. Changes shared live! Use sidebar to save/load.")
        current_file_basename = st.session_state.get('current_world_file', None)
        if current_file_basename:
             full_path = os.path.join(SAVED_WORLDS_DIR, current_file_basename)
             if os.path.exists(full_path): parsed = parse_world_filename(full_path); st.info(f"🌠 Viewing: **{parsed['name']}** (`{current_file_basename}`)")
             else: st.warning(f"⚠️ Loaded file '{current_file_basename}' missing."); st.session_state.current_world_file = None
        else: st.info("☁️ Live State Active (Save to persist)")

        # 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
            try: # Get WS URL (Best effort)
                 from streamlit.web.server.server import Server
                 session_info = Server.get_current()._get_session_info(st.runtime.scriptrunner.get_script_run_ctx().session_id)
                 host_attr = getattr(session_info.ws.stream.request, 'host', None) or getattr(getattr(session_info, 'client', None), 'request', None)
                 if host_attr: server_host = host_attr.host.split(':')[0]; ws_url = f"ws://{server_host}:8765"
                 else: raise AttributeError("Host attribute not found")
            except Exception as e: print(f"⚠️ WS URL detection failed ({e}), using localhost.")

            # Inject only necessary state for JS init
            js_injection_script = f"""<script>
    window.USERNAME = {json.dumps(st.session_state.username)};
    window.WEBSOCKET_URL = {json.dumps(ws_url)}; // Needed by JS to connect
    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)};
    // Initial world state is sent via WebSocket 'initial_state' message now
    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"💬 Whispers in the Void")
        chat_history_list = st.session_state.get('chat_history', [])
        if not chat_history_list: chat_history_list = asyncio.run(load_chat_history())
        chat_container = st.container(height=500)
        with chat_container:
             if chat_history_list: st.markdown("----\n".join(reversed(chat_history_list[-50:])))
             else: st.caption("🌬️ Silence reigns...")

        def clear_chat_input_callback(): st.session_state.message_input = ""
        message_value = st.text_input("Your Message:", key="message_input", label_visibility="collapsed")
        send_button_clicked = st.button("✉️ Send Message", key="send_chat_button", on_click=clear_chat_input_callback)

        if send_button_clicked:
            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
                voice = st.session_state.get('tts_voice', DEFAULT_TTS_VOICE)
                ws_message = json.dumps({"type": "chat_message", "payload": {"username": st.session_state.username, "message": message_to_send, "voice": voice}})
                # Use run_async for background tasks
                run_async(broadcast_message, ws_message) # Broadcast Chat via WS
                run_async(save_chat_entry, st.session_state.username, message_to_send, voice) # Save async
                add_action_log(f"Sent chat: {message_to_send[:20]}...", emoji="💬")
                # Rerun is handled implicitly by button + on_click
            elif send_button_clicked: st.toast("Message empty or same as last.")

    # --- PDF Tab ---
    with tab_pdf:
         st.header("📚 Tome Translator")
         st.caption("🔊 Give voice to the silent knowledge within PDF scrolls.")
         pdf_file = st.file_uploader("Upload PDF Scroll", type="pdf", key="pdf_upload")
         max_pages = st.slider('Max Pages to Animate', 1, 50, 10, key="pdf_pages")
         if pdf_file:
             if st.button("🎙️ Animate PDF to Audio", key="process_pdf_button"):
                 with st.spinner("✨ Transcribing ancient glyphs to sound..."):
                     process_pdf_tab(pdf_file, max_pages, st.session_state.tts_voice)

    # --- Files & Settings Tab ---
    with tab_files:
        st.header("📂 Archives & Settings")
        st.caption("⚙️ Manage saved scrolls and application settings.")

        st.subheader("💾 World Scroll Management")
        current_file_basename = st.session_state.get('current_world_file', None)

        # Save Current Version Button
        if current_file_basename:
            full_path = os.path.join(SAVED_WORLDS_DIR, current_file_basename)
            save_label = f"Save Changes to '{current_file_basename}'"
            if os.path.exists(full_path): parsed = parse_world_filename(full_path); save_label = f"💾 Save Changes to '{parsed['name']}'"
            if st.button(save_label, key="save_current_world_files", help=f"Overwrite '{current_file_basename}'"):
                 if not os.path.exists(full_path): st.error(f"❌ Cannot save, file missing.")
                 else:
                     with st.spinner(f"Saving changes to {current_file_basename}..."):
                         # Save the current state from the cached resource
                         if save_world_state_to_md(current_file_basename):
                             st.success("✅ Current world saved!"); add_action_log(f"Saved world: {current_file_basename}", emoji="💾")
                         else: st.error("❌ Failed to save world state.")
        else: st.info("➡️ Load a world from the sidebar to enable 'Save Changes'.")

        # Save As New Version Section
        st.subheader("✨ Save As New Scroll")
        new_name_files = st.text_input("New Scroll Name:", key="new_world_name_files_tab", value=st.session_state.get('new_world_name', 'MyDreamscape'))
        if st.button("💾 Save Current View as New Scroll", key="save_new_version_files"):
             if new_name_files.strip():
                 with st.spinner(f"Saving new version '{new_name_files}'..."):
                    new_filename_base = generate_world_save_filename(st.session_state.username, new_name_files)
                    # Save the current state from the cached resource to a NEW file
                    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; st.session_state.new_world_name = "MyDreamscape";
                         add_action_log(f"Saved new world: {new_filename_base}", emoji="✨")
                         st.rerun()
                    else: st.error("❌ Failed to save new version.")
             else: st.warning("⚠️ Please enter a name.")

        # Server Status
        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:
             with clients_lock: client_count = len(connected_clients)
             st.metric("🔗 Connected Clients", client_count)

        # File Deletion
        st.subheader("🗑️ Archive Maintenance")
        st.caption("🧹 Cleanse the old to make way for the new.")
        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 = []; add_action_log("Cleared Chats", emoji="🧹"); 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 = {}; add_action_log("Cleared Audio", emoji="🧹"); 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")]); st.session_state.current_world_file = None; add_action_log("Cleared Worlds", emoji="🧹"); st.rerun()
        with col_del4:
            if st.button("🗑️ All Gen", key="del_all_gen"): 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"), os.path.join(MEDIA_DIR, "*.zip")]); st.session_state.chat_history = []; st.session_state.audio_cache = {}; st.session_state.current_world_file = None; add_action_log("Cleared All Generated", emoji="🔥"); st.rerun()

        # Download Archives
        st.subheader("📦 Download Archives")
        st.caption("Bundle your creations for safekeeping or sharing.")
        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, f"{WORLD_STATE_FILE_MD_PREFIX}*.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")
        zip_files = sorted(glob.glob(os.path.join(MEDIA_DIR,"*.zip")), key=os.path.getmtime, reverse=True)
        if zip_files:
            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 archives found.")


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

def initialize_app():
    """Handles session init, server start, and ensures world state resource is accessed."""
    init_session_state()
    # Load/Assign username
    if not st.session_state.username:
        loaded_user = load_username()
        if loaded_user and loaded_user in FUN_USERNAMES: st.session_state.username = loaded_user; st.session_state.tts_voice = FUN_USERNAMES[loaded_user]
        else: st.session_state.username = random.choice(list(FUN_USERNAMES.keys())) if FUN_USERNAMES else "User"; st.session_state.tts_voice = FUN_USERNAMES.get(st.session_state.username, DEFAULT_TTS_VOICE); save_username(st.session_state.username)

    # Ensure WebSocket server thread is running
    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): st.session_state.server_running_flag = True

    # Trigger the cached resource initialization/retrieval
    try:
        manager = get_world_state_manager()
        # Set initial current_world_file if needed (based on what cache loaded)
        if st.session_state.get('current_world_file') is None:
            if manager["state"]: # If the cache loaded state from a file
                saved_worlds = get_saved_worlds()
                if saved_worlds:
                    st.session_state.current_world_file = os.path.basename(saved_worlds[0]['filename'])
                    print(f"🐍 Set initial session 'current_world_file' to: {st.session_state.current_world_file}")
    except Exception as e:
        st.error(f"❌ Fatal error initializing world state manager: {e}"); st.exception(e); st.stop()


if __name__ == "__main__":
    initialize_app()
    render_sidebar()
    render_main_content()