# app.py (Refactored & Consolidated)
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

# ==============================================================================
# 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": "📄"}

# API Keys (Load from .env or secrets)
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', ""))

# Mapping Emojis to Primitive Types
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"
}

# ==============================================================================
# Global State & Locks
# ==============================================================================

# Thread lock for accessing shared world state
world_objects_lock = threading.Lock()
# In-memory world state {obj_id: data} - Use defaultdict for convenience
world_objects = defaultdict(dict)
# Set of active WebSocket client IDs
connected_clients = set()

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

def get_current_time_str(tz='UTC'):
    """Gets formatted timestamp string in specified timezone (default UTC)."""
    try:
        timezone = pytz.timezone(tz)
        # Get current datetime localized to the specified timezone
        now_aware = datetime.now(timezone)
    except pytz.UnknownTimeZoneError:
        # Fallback to UTC if timezone is unknown
        now_aware = datetime.now(pytz.utc)
    except Exception as e:
         # General fallback if timezone localization fails
         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."""
    try:
        loop = asyncio.get_running_loop()
        # Schedule as task if loop is running
        return loop.create_task(async_func(*args, **kwargs))
    except RuntimeError: # No running loop
        # Run in a new loop (blocks until completion)
        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 # Indicate error or failure
    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)
    # Ensure prefix and suffix are correct
    if basename.startswith(WORLD_STATE_FILE_MD_PREFIX) and basename.endswith(".md"):
        parts = basename[len(WORLD_STATE_FILE_MD_PREFIX):-3].split('_') # Remove prefix and suffix before split
        if len(parts) >= 3: # Expecting Name_Timestamp_Hash
            timestamp_str = parts[-2]
            name = " ".join(parts[:-2]) # Join potentially multiple parts for name
            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 since generated with UTC
            except (ValueError, pytz.exceptions.AmbiguousTimeError, pytz.exceptions.NonExistentTimeError):
                dt_obj = None # Parsing failed or timezone issue
            return {"name": name or "Untitled", "timestamp": timestamp_str, "dt": dt_obj, "filename": filename}

    # Fallback for unknown format or parsing failure
    dt_fallback = None
    try:
         mtime = os.path.getmtime(filename)
         dt_fallback = datetime.fromtimestamp(mtime, tz=pytz.utc)
    except Exception: pass # Ignore errors getting mtime
    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 {len(world_objects)} objects to: {save_path}...")
    success = False
    with world_objects_lock:
        world_data_dict = dict(world_objects) # Create copy inside lock
        print(f"Saving {len(world_data_dict)} objects...")
        parsed_info = parse_world_filename(target_filename_base) # Use original filename for info
        timestamp_save = get_current_time_str()
        md_content = f"""# World State: {parsed_info['name']}
* **File Saved:** {timestamp_save} (UTC)
* **Original Timestamp:** {parsed_info['timestamp']}
* **Objects:** {len(world_data_dict)}

```json
{json.dumps(world_data_dict, indent=2)}
```"""
        try:
            ensure_dir(SAVED_WORLDS_DIR) # Ensure dir exists
            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()
        json_match = re.search(r"```json\s*(\{.*?\})\s*```", content, re.DOTALL | 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() # Clear previous state
             for k, v in world_data_dict.items(): world_objects[str(k)] = v # Update with loaded data
             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 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]
        # Sort by datetime object if available (newest first)
        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) # Remove markdown links, keep text
    text = re.sub(r'[#*_`!]', '', text) # Remove some markdown chars
    text = ' '.join(text.split()) # Normalize whitespace
    return text[:250] or "No text"

def create_file(content, username, file_type="md", save_path=None):
    if not save_path:
         # Generate filename if specific path not given
         filename = generate_filename(content, username, file_type)
         save_path = os.path.join(MEDIA_DIR, filename) # Save to base dir by default
    ensure_dir(os.path.dirname(save_path)) # Ensure directory exists
    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)
            link_html = f'<a href="data:{mime_types.get(file_type, "application/octet-stream")};base64,{b64}" download="{basename}">{FILE_EMOJIS.get(file_type, "📄")} DL</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 based on text hash and voice
    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: {e}")
    else: st.warning(f"Audio file not found: {os.path.basename(file_path) if file_path else 'N/A'}")

# --- 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)
    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."""
    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

# --- 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):
    protected = [STATE_FILE, "app.py", "index.html", "requirements.txt", "README.md"]
    # Dynamically protect all currently saved world files unless explicitly targeted
    try:
        current_worlds = [os.path.basename(w['filename']) for w in get_saved_worlds()]
        protected.extend(current_worlds)
    except Exception: pass # Ignore if world listing fails
    if exclude_files: protected.extend(exclude_files)

    deleted_count = 0; errors = 0
    for pattern in file_patterns:
         pattern_path = os.path.join(MEDIA_DIR, pattern) # Assume relative to app dir
         try:
              files_to_delete = glob.glob(pattern_path)
              if not files_to_delete: 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}")
         except Exception as glob_e: print(f"Err matching {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 (Ctrl+V)", 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
      return pasted_img if submit_button and pasted_img else None


# --- PDF Processing ---
# Note: Depends on PyPDF2 and potentially AudioProcessor class if generating audio
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=f"{self.cache_dir}/{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):
      st.subheader("PDF Processing")
      if pdf_file is None:
           st.info("Upload a PDF file to begin.")
           return
      audio_processor = AudioProcessor() # Instance for this run
      try:
          reader=PdfReader(pdf_file); 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):
              # Runs async audio generation in a separate thread
              async def run_async_audio(): return await audio_processor.create_audio(page_text, voice)
              try: audio_path = asyncio.run(run_async_audio());
                   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):
               text=reader.pages[i].extract_text();
               if text: 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]"

          # Display results as they become available (or after join)
          st.progress(0) # Placeholder for progress
          # Wait for threads and display - consider using st.empty() for updates
          for thread in page_threads: thread.join()

          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 pending.")

      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);
    username = st.session_state.get('username', f"User_{client_id[:4]}") # Get username for this session
    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", {});
                # Use username from payload if provided, otherwise session username
                sender_username = payload.get("username", username)

                if msg_type == "chat_message":
                    chat_text = payload.get('message', ''); voice = payload.get('voice', FUN_USERNAMES.get(sender_username, "en-US-AriaNeural"));
                    # Schedule save/TTS, but broadcast immediately
                    run_async(save_chat_entry, sender_username, chat_text, voice) # Fire and forget save
                    await broadcast_message(message, exclude_id=client_id) # Forward original msg

                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") # Optionally include 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)

                # Add more handlers here

            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() # Keep running
    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():
        loop = asyncio.new_event_loop(); asyncio.set_event_loop(loop)
        try: loop.run_until_complete(run_websocket_server())
        finally: 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
# ==============================================================================

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

        # World Selector
        saved_worlds = get_saved_worlds()
        # Create display text mapping filename to formatted string
        world_options = {w['filename']: f"{w['name']} ({w['timestamp']})" for w in saved_worlds}
        # Get the currently selected filename (basename) from session state
        current_selection_basename = st.session_state.get('current_world_file', None)

        # Prepare options for radio button: None first, then filenames
        radio_options_list = [None] + [w['filename'] for w in saved_worlds] # Store full path initially? No, use basename.
        radio_options_basenames = [None] + [os.path.basename(w['filename']) for w in saved_worlds]

        # Find the index of the current selection in the options list
        current_radio_index = 0 # Default to "Live State"
        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 # Fallback if filename somehow not in list

        # Display radio buttons
        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.get(os.path.join(SAVED_WORLDS_DIR, x), x), # Format using full path to get name/time? complex
             key="world_selector_radio"
         )

        # Handle selection change
        if selected_basename != current_selection_basename:
             st.session_state.current_world_file = selected_basename # Store 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 # Reset on failure
             else:
                  print("Switched to live state.")
                  # Optionally clear world or just stop tracking file? Stop tracking.
                  # Maybe broadcast current live state if switching TO live?
                  # run_async(broadcast_world_update) # Broadcast current live state
                  st.toast("Switched to Live State.")
             st.rerun()

        # Display download links
        st.caption("Download:")
        cols = st.columns([3, 1]) # Columns for name and download button
        with cols[0]: st.write("**Name**")
        with cols[1]: st.write("**Link**")
        for world_info in saved_worlds:
            f_basename = os.path.basename(world_info['filename'])
            f_fullpath = world_info['filename'] # Full path needed for reading file
            display_name = world_info.get('name', f_basename)
            timestamp = world_info.get('timestamp', 'N/A')
            col1, col2 = st.columns([3, 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("---")

        # Save New Version moved to Files Tab

        # 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
                    run_async(lambda: streamlit_js_eval(f"updateSelectedObjectType({json.dumps(name)});", key=f"update_tool_js_{name}")) # Fire-and-forget JS update
                    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
        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)
            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 = st.session_state.get('current_world_file', None)
        if current_file: parsed = parse_world_filename(current_file); st.info(f"Current World: **{parsed['name']}** (`{os.path.basename(current_file)}`)")
        else: st.info("Live State Active (Unsaved changes will be lost unless saved as new version)")

        # Embed HTML Component
        html_file_path = 'index.html'
        try:
            with open(html_file_path, 'r', encoding='utf-8') as f: html_template = f.read()
            try: # Get WS URL
                 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: ws_url = "ws://localhost:8765"

            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")
        chat_history = run_async(load_chat_history).result() if 'chat_history' not in st.session_state else st.session_state.chat_history # Load sync if needed
        chat_container = st.container(height=500)
        with chat_container:
             if chat_history: st.markdown("----\n".join(reversed(chat_history[-50:])))
             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
                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}})
                run_async(broadcast_message, ws_message) # Send via WS
                run_async(save_chat_entry, st.session_state.username, message_to_send, voice) # Save locally
                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") # Limit pages
         if pdf_file:
              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 = st.session_state.get('current_world_file', None)

        # Save Current Version Button
        if current_file:
             parsed = parse_world_filename(current_file); save_label = f"Save Changes to '{parsed['name']}'"
             if st.button(save_label, key="save_current_world", help=f"Overwrite '{os.path.basename(current_file)}' with current live state."):
                 with st.spinner(f"Overwriting {os.path.basename(current_file)}..."):
                     if save_world_state_to_md(current_file): st.success("Current world version saved!")
                     else: st.error("Failed to save world state.")
         else:
             st.info("Load a world version from the sidebar or use 'Save New Version' below to save the current live state.")

        # 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") # Use different key from sidebar one if kept there
        if st.button("💾 Save Live State as New Version", key="save_new_version_files"):
             if new_name_files.strip():
                  new_filename = generate_world_save_filename(new_name_files) # Generates basename
                  with st.spinner(f"Saving new version '{new_name_files}'..."):
                      if save_world_state_to_md(new_filename): # Pass basename
                          st.success(f"Saved as {new_filename}")
                          st.session_state.current_world_file = new_filename # Switch to new file automatically
                          st.session_state.new_world_name_files = "" # Reset input
                          st.rerun()
                      else: st.error("Failed to save new version.")
             else: st.warning("Please enter a name for the new world version.")

        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", help="Deletes world save files (.md) in saved_worlds"): delete_files([os.path.join(SAVED_WORLDS_DIR, f"{WORLD_STATE_FILE_MD_PREFIX}*.md")]); st.session_state.current_world_file = None; st.rerun()
        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("*.zip"), key=os.path.getmtime, reverse=True)
        if 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."""
    if not st.session_state.get('initial_world_state_loaded', False):
        print("Performing initial world load...")
        saved_worlds = get_saved_worlds()
        if saved_worlds:
             latest_world_file_basename = os.path.basename(saved_worlds[0]['filename'])
             print(f"Loading most recent world on startup: {latest_world_file_basename}")
             load_world_state_from_md(latest_world_file_basename) # This updates global state and sets session state 'current_world_file'
         else:
              print("No saved worlds found, starting with empty state.")
              with world_objects_lock: world_objects.clear() # Ensure empty state
              st.session_state.current_world_file = None # Ensure no file is marked as loaded
         st.session_state.initial_world_state_loaded = True
         print("Initial world load complete.")


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

    # 2. Start WebSocket server thread if needed
    # Use server_running_flag to prevent multiple start attempts within one session
    if not st.session_state.get('server_running_flag', False):
         if 'server_task' not in st.session_state or not st.session_state.server_task.is_alive():
              start_websocket_server_thread()
         else:
              # If task exists but flag is false, maybe update flag?
              if st.session_state.server_task.is_alive():
                  st.session_state.server_running_flag = True
                  print("Corrected server_running_flag based on alive thread.")


    # 3. Load initial world state from disk if needed
    initialize_world()

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