app.py

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
from streamlit_js_eval import streamlit_js_eval
from PIL import Image

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

nest_asyncio.apply()

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

icons = '🤖🏗️🗣️💾'
Site_Name = '🤖🏗️ Shared World Builder 🗣️'
START_ROOM = "World Lobby 🌍"
MEDIA_DIR = "."
STATE_FILE = "user_state.txt"

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_CACHE_DIR = "audio_cache"
AUDIO_DIR = "audio_logs"
SAVED_WORLDS_DIR = "saved_worlds"
HISTORY_LOG_DIR = "history_logs"
PLOT_WIDTH = 50.0
PLOT_DEPTH = 50.0
WORLD_STATE_FILE_MD_PREFIX = "🌍_"

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

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"
}

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

world_objects_lock = threading.Lock()
world_objects = defaultdict(dict)
connected_clients = set()

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

def get_current_time_str(tz='UTC'):
    try:
        timezone = pytz.timezone(tz)
        now_aware = datetime.now(timezone)
    except pytz.UnknownTimeZoneError:
        now_aware = datetime.now(pytz.utc)
    return now_aware.strftime('%Y%m%d_%H%M%S')

def clean_filename_part(text, max_len=30):
    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):
    try:
        loop = asyncio.get_running_loop()
        return loop.create_task(async_func(*args, **kwargs))
    except RuntimeError:
        print(f"Warning: Running async func {async_func.__name__} in new event loop.")
        try:
            return asyncio.run(async_func(*args, **kwargs))
        except Exception as e:
            print(f"Error running async func {async_func.__name__} in new loop: {e}")
            return None
    except Exception as e:
        print(f"Error scheduling async task {async_func.__name__}: {e}")
        return None

def ensure_dir(dir_path):
    os.makedirs(dir_path, exist_ok=True)

def generate_filename(content, username, extension):
    timestamp = get_current_time_str()
    content_hash = hashlib.md5(content[:150].encode()).hexdigest()[:6]
    clean_username = clean_filename_part(username)
    return f"{clean_username}_{timestamp}_{content_hash}.{extension}"

def format_timestamp_prefix(prefix):
    return f"{prefix}_{get_current_time_str()}"

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

def generate_world_save_filename(name="World"):
    timestamp = get_current_time_str()
    clean_name = clean_filename_part(name)
    rand_hash = hashlib.md5(str(time.time()).encode() + name.encode()).hexdigest()[:6]
    return f"{WORLD_STATE_FILE_MD_PREFIX}{clean_name}_{timestamp}_{rand_hash}.md"

def parse_world_filename(filename):
    basename = os.path.basename(filename)
    if basename.startswith(WORLD_STATE_FILE_MD_PREFIX) and basename.endswith(".md"):
        core_name = basename[len(WORLD_STATE_FILE_MD_PREFIX):-3]
        parts = core_name.split('_')
        if len(parts) >= 3:
            timestamp_str = parts[-2]
            name_parts = parts[:-2]
            name = "_".join(name_parts) if name_parts else "Untitled"
            dt_obj = None
            try:
                dt_obj = datetime.strptime(timestamp_str, '%Y%m%d_%H%M%S')
                dt_obj = pytz.utc.localize(dt_obj)
            except (ValueError, pytz.exceptions.AmbiguousTimeError, pytz.exceptions.NonExistentTimeError):
                dt_obj = None
            return {"name": name.replace('_', ' '), "timestamp": timestamp_str, "dt": dt_obj, "filename": filename}
    dt_fallback = None
    try:
        mtime = os.path.getmtime(filename)
        dt_fallback = datetime.fromtimestamp(mtime, tz=pytz.utc)
    except Exception:
        pass
    return {"name": basename.replace('.md',''), "timestamp": "Unknown", "dt": dt_fallback, "filename": filename}

def save_world_state_to_md(target_filename_base):
    global world_objects
    save_path = os.path.join(SAVED_WORLDS_DIR, target_filename_base)
    print(f"Acquiring lock to save world state to: {save_path}...")
    success = False
    with world_objects_lock:
        world_data_dict = dict(world_objects)
        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']}
* **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
            username = st.session_state.get('username', 'Anonymous')
            clean_username = clean_filename_part(username)
            player_filename_base = f"{clean_username}_{target_filename_base}"
            player_save_path = os.path.join(SAVED_WORLDS_DIR, player_filename_base)
            with open(player_save_path, 'w', encoding='utf-8') as f:
                f.write(md_content)
            print(f"Player-specific world state saved to {player_filename_base}")
        except Exception as e:
            print(f"Error saving world state to {save_path}: {e}")
    return success

def load_world_state_from_md(filename_base):
    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\S]*?\})\s*```", content, re.IGNORECASE)
        if not json_match:
            st.error(f"Could not find valid JSON block in {filename_base}")
            return False
        world_data_dict = json.loads(json_match.group(1))
        print(f"Acquiring lock to update world state from {filename_base}...")
        with world_objects_lock:
            world_objects.clear()
            for k, v in world_data_dict.items():
                world_objects[str(k)] = v
            loaded_count = len(world_objects)
        print(f"Loaded {loaded_count} objects from {filename_base}. Lock released.")
        st.session_state.current_world_file = filename_base
        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():
    try:
        ensure_dir(SAVED_WORLDS_DIR)
        world_files = glob.glob(os.path.join(SAVED_WORLDS_DIR, "*.md"))
        parsed_worlds = [parse_world_filename(f) for f in world_files]
        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 []

def log_action(username, action_type, data):
    timestamp = get_current_time_str()
    log_entry = {
        "timestamp": timestamp,
        "username": username,
        "action": action_type,
        "data": data
    }
    shared_log_file = os.path.join(HISTORY_LOG_DIR, "shared_history.jsonl")
    try:
        with open(shared_log_file, 'a', encoding='utf-8') as f:
            f.write(json.dumps(log_entry) + '\n')
    except Exception as e:
        print(f"Error writing to shared history log: {e}")
    clean_username = clean_filename_part(username)
    player_log_file = os.path.join(HISTORY_LOG_DIR, f"{clean_username}_history.jsonl")
    try:
        with open(player_log_file, 'a', encoding='utf-8') as f:
            f.write(json.dumps(log_entry) + '\n')
    except Exception as e:
        print(f"Error writing to player history log {player_log_file}: {e}")

# ==============================================================================
# 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():
    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, '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
    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 = {}
    if 'username' not in st.session_state or not st.session_state.username:
        st.session_state.username = random.choice(list(FUN_USERNAMES.keys()))
        save_username(st.session_state.username)

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

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

def create_file(content, username, file_type="md", save_path=None):
    if not save_path:
        filename = generate_filename(content, username, file_type)
        save_path = os.path.join(MEDIA_DIR, filename)
    ensure_dir(os.path.dirname(save_path))
    try:
        with open(save_path, 'w', encoding='utf-8') as f:
            f.write(content)
        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"):
    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):
    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):
    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):
    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():
    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"):
    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"]
    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
                elif os.path.isdir(f_path):
                    print(f"Skipping directory: {f_path}")
        except Exception as glob_e:
            print(f"Error matching pattern {pattern}: {glob_e}")
            errors += 1
    msg = f"Deleted {deleted_count} files."
    if errors > 0:
        msg += f" Encountered {errors} errors."
    st.warning(msg) if errors > 0 else st.success(msg) if deleted_count > 0 else st.info("No matching files found to delete.")
    st.session_state['download_link_cache'] = {}
    st.session_state['audio_cache'] = {}

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
    paste_input = st.text_area("Paste Image Data Here", key="paste_input_area", height=50, value="")
    if st.button("Process Pasted Image 📋", key="paste_form_button"):
        if 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 = ""
        else:
            st.warning("No valid image data pasted.")
            st.session_state.paste_image_base64 = ""
    return pasted_img

class AudioProcessor:
    def __init__(self):
        self.cache_dir = AUDIO_CACHE_DIR
        ensure_dir(self.cache_dir)
        self.metadata = json.load(open(f"{self.cache_dir}/metadata.json", 'r')) if os.path.exists(f"{self.cache_dir}/metadata.json") else {}
    def _save_metadata(self):
        try:
            with open(f"{self.cache_dir}/metadata.json", 'w') as f:
                json.dump(self.metadata, f, indent=2)
        except Exception as e:
            print(f"Failed metadata save: {e}")
    async def create_audio(self, text, voice='en-US-AriaNeural'):
        cache_key = hashlib.md5(f"{text[:150]}:{voice}".encode()).hexdigest()
        cache_path = os.path.join(self.cache_dir, f"{cache_key}.mp3")
        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 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 = min(len(reader.pages), max_pages)
        st.write(f"Processing first {total_pages} pages of '{pdf_file.name}'...")
        texts, audios = {}, {}
        page_threads = []
        results_lock = threading.Lock()

        def process_page_sync(page_num, page_text):
            async def run_async_audio():
                return await audio_processor.create_audio(page_text, voice)
            try:
                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}")

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

        progress_bar = st.progress(0.0, text="Processing pages...")
        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.")
                break
            time.sleep(0.5)
        progress_bar.progress(1.0, text="Processing complete.")

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

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

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

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

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

async def send_safely(websocket, message, client_id):
    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):
    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)

async def broadcast_world_update():
    with world_objects_lock:
        current_state_payload = dict(world_objects)
    update_msg = json.dumps({"type": "initial_state", "payload": current_state_payload})
    print(f"Broadcasting full world update ({len(current_state_payload)} objects)...")
    await broadcast_message(update_msg)

async def websocket_handler(websocket, path):
    await register_client(websocket)
    client_id = str(websocket.id)
    username = st.session_state.get('username', f"User_{client_id[:4]}")
    try:
        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:
        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)

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

                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)
                        current_file = st.session_state.get('current_world_file')
                        if current_file:
                            run_async(save_world_state_to_md, current_file)
                        else:
                            new_filename = generate_world_save_filename(f"AutoSave_{sender_username}")
                            if run_async(save_world_state_to_md, new_filename):
                                st.session_state.current_world_file = new_filename
                        log_action(sender_username, "place", obj_data)

                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)
                            current_file = st.session_state.get('current_world_file')
                            if current_file:
                                run_async(save_world_state_to_md, current_file)
                            else:
                                new_filename = generate_world_save_filename(f"AutoSave_{sender_username}")
                                if run_async(save_world_state_to_md, new_filename):
                                    st.session_state.current_world_file = new_filename
                            log_action(sender_username, "delete", {"obj_id": obj_id})

                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)

                elif msg_type == "tool_selected":
                    tool_type = payload.get("tool_type")
                    if tool_type in PRIMITIVE_MAP.values() or tool_type == "None":
                        st.session_state.selected_object = tool_type
                        print(f"Tool updated to {tool_type} for {sender_username}")

            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():
    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():
    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 = 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)
                for task in tasks:
                    task.cancel()
                try:
                    loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=True))
                except asyncio.CancelledError:
                    pass
                finally:
                    loop.close()
                    print("Server thread loop closed.")
            else:
                print("Server thread loop already closed or None.")
    st.session_state.server_task = threading.Thread(target=run_loop, daemon=True)
    st.session_state.server_task.start()
    time.sleep(1.5)
    if not st.session_state.server_task.is_alive():
        print("### Server thread failed to stay alive!")

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

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

        saved_worlds = get_saved_worlds()
        world_options_display = {
            os.path.basename(w['filename']): (
                f"{w['name']} ({w['timestamp']})" if not w['filename'].startswith(clean_filename_part(st.session_state.get('username', '')))
                else f"{w['name']} ({w['timestamp']}) [Your Save]"
            ) for w in saved_worlds
        }
        radio_options_basenames = [None] + [os.path.basename(w['filename']) for w in saved_worlds]
        current_selection_basename = st.session_state.get('current_world_file', None)
        current_radio_index = 0
        if current_selection_basename and current_selection_basename in radio_options_basenames:
            try:
                current_radio_index = radio_options_basenames.index(current_selection_basename)
            except ValueError:
                current_radio_index = 0

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

        if selected_basename != current_selection_basename:
            st.session_state.current_world_file = selected_basename
            if selected_basename:
                with st.spinner(f"Loading {selected_basename}..."):
                    if load_world_state_from_md(selected_basename):
                        run_async(broadcast_world_update)
                        st.toast("World loaded!", icon="✅")
                    else:
                        st.error("Failed to load world.")
                        st.session_state.current_world_file = None
            else:
                print("Switched to live state.")
                st.toast("Switched to Live State.")
            st.rerun()

        st.caption("Download:")
        cols = st.columns([4, 1])
        with cols[0]:
            st.write("**Name** (Timestamp)")
        with cols[1]:
            st.write("**DL**")
        display_limit = 10
        for i, world_info in enumerate(saved_worlds):
            if i >= display_limit and len(saved_worlds) > display_limit + 1:
                with st.expander(f"Show {len(saved_worlds)-display_limit} more..."):
                    for world_info_more in saved_worlds[display_limit:]:
                        f_basename = os.path.basename(world_info_more['filename'])
                        f_fullpath = os.path.join(SAVED_WORLDS_DIR, f_basename)
                        display_name = world_info_more.get('name', f_basename)
                        timestamp = world_info_more.get('timestamp', 'N/A')
                        colA, colB = st.columns([4, 1])
                        with colA:
                            st.write(f"<small>{display_name} ({timestamp})</small>", unsafe_allow_html=True)
                        with colB:
                            st.markdown(get_download_link(f_fullpath, "md"), unsafe_allow_html=True)
                break
            f_basename = os.path.basename(world_info['filename'])
            f_fullpath = os.path.join(SAVED_WORLDS_DIR, f_basename)
            display_name = world_info.get('name', f_basename)
            timestamp = world_info.get('timestamp', 'N/A')
            col1, col2 = st.columns([4, 1])
            with col1:
                st.write(f"<small>{display_name} ({timestamp})</small>", unsafe_allow_html=True)
            with col2:
                st.markdown(get_download_link(f_fullpath, "md"), unsafe_allow_html=True)

        st.markdown("---")
        st.header("🏗️ Build Tools")
        st.caption("Select an object to place.")
        # Hidden selectbox to sync tool state
        tool_options = ['None'] + list(PRIMITIVE_MAP.values())
        selected_tool = st.selectbox("Selected Tool", options=tool_options, index=tool_options.index(st.session_state.get('selected_object', 'None')), key="tool_select", label_visibility="collapsed")
        if selected_tool != st.session_state.selected_object:
            st.session_state.selected_object = selected_tool
            run_async(lambda: streamlit_js_eval(
                f"updateSelectedObjectType({json.dumps(selected_tool)});",
                key=f"update_tool_js_{selected_tool}"
            ))

        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):
                # Update tool via JavaScript without rerun
                run_async(lambda name_arg=name: streamlit_js_eval(
                    f"""
                    updateSelectedObjectType('{name_arg}');
                    websocket.send(JSON.stringify({{
                        type: 'tool_selected',
                        payload: {{ tool_type: '{name_arg}', username: window.USERNAME }}
                    }}));
                    """,
                    key=f"update_tool_js_{name_arg}"
                ))
            col_idx += 1
        st.markdown("---")
        if st.button("🚫 Clear Tool", key="clear_tool", use_container_width=True):
            run_async(lambda: streamlit_js_eval(
                """
                updateSelectedObjectType('None');
                websocket.send(JSON.stringify({
                    type: 'tool_selected',
                    payload: { tool_type: 'None', username: window.USERNAME }
                }));
                """,
                key="update_tool_js_none"
            ))

        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():
    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"])

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

        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"
            try:
                from streamlit.web.server.server import Server
                session_info = Server.get_current()._get_session_info(st.runtime.scriptrunner.get_script_run_ctx().session_id)
                server_host = session_info.ws.stream.request.host.split(':')[0]
                ws_url = f"ws://{server_host}:8765"
            except Exception as e:
                print(f"WS URL detection failed ({e}), using localhost.")
            js_injection_script = f"""<script>
    window.USERNAME = {json.dumps(st.session_state.username)};
    window.WEBSOCKET_URL = {json.dumps(ws_url)};
    window.SELECTED_OBJECT_TYPE = {json.dumps(st.session_state.selected_object)};
    window.PLOT_WIDTH = {json.dumps(PLOT_WIDTH)};
    window.PLOT_DEPTH = {json.dumps(PLOT_DEPTH)};
    console.log("Streamlit State Injected:", {{ username: window.USERNAME, websocketUrl: window.WEBSOCKET_URL, selectedObject: window.SELECTED_OBJECT_TYPE }});
</script>"""
            html_content_with_state = html_template.replace('</head>', js_injection_script + '\n</head>', 1)
            components.html(html_content_with_state, height=700, scrolling=False)
        except FileNotFoundError:
            st.error(f"CRITICAL ERROR: Could not find '{html_file_path}'.")
        except Exception as e:
            st.error(f"Error loading 3D component: {e}")
            st.exception(e)

    with tab_chat:
        st.header(f"{START_ROOM} Chat")
        chat_history_task = run_async(load_chat_history)
        chat_history = chat_history_task.result() if chat_history_task else st.session_state.chat_history
        chat_container = st.container(height=500)
        with chat_container:
            if chat_history:
                st.markdown("----\n".join(reversed(chat_history[-50:])))
            else:
                st.caption("No chat messages yet.")

        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)
                run_async(save_chat_entry, st.session_state.username, message_to_send, voice)
                st.session_state.message_input = ""
                st.rerun()
            elif send_button_clicked:
                st.toast("Message empty or same as last.")
        st.checkbox("Autosend Chat", key="autosend")

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

    with tab_files:
        st.header("📂 Files & Settings")
        st.subheader("💾 World State Management")
        current_file_basename = st.session_state.get('current_world_file', None)
        if current_file_basename:
            parsed = parse_world_filename(os.path.join(SAVED_WORLDS_DIR, current_file_basename))
            save_label = f"Save Changes to '{parsed['name']}'"
            if st.button(save_label, key="save_current_world", help=f"Overwrite '{current_file_basename}'"):
                with st.spinner(f"Overwriting {current_file_basename}..."):
                    if save_world_state_to_md(current_file_basename):
                        st.success("Current world saved!")
                    else:
                        st.error("Failed to save world state.")
        else:
            st.info("Load a world from the sidebar to enable saving changes to it.")

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

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

        st.subheader("🗑️ Delete Files")
        st.warning("Deletion is permanent!", icon="⚠️")
        col_del1, col_del2, col_del3, col_del4 = st.columns(4)
        with col_del1:
            if st.button("🗑️ Chats", key="del_chat_md"):
                delete_files([os.path.join(CHAT_DIR, "*.md")])
                st.session_state.chat_history = []
                st.rerun()
        with col_del2:
            if st.button("🗑️ Audio", key="del_audio_mp3"):
                delete_files([os.path.join(AUDIO_DIR, "*.mp3"), os.path.join(AUDIO_CACHE_DIR, "*.mp3")])
                st.session_state.audio_cache = {}
                st.rerun()
        with col_del3:
            if st.button("🗑️ Worlds", key="del_worlds_md"):
                delete_files([os.path.join(SAVED_WORLDS_DIR, f"{WORLD_STATE_FILE_MD_PREFIX}*.md")])
                st.session_state.current_world_file = None
                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"), "*.zip"])
                st.session_state.chat_history = []
                st.session_state.audio_cache = {}
                st.session_state.current_world_file = None
                st.rerun()

        st.subheader("📦 Download Archives")
        zip_files = sorted(glob.glob(os.path.join(MEDIA_DIR,"*.zip")), key=os.path.getmtime, reverse=True)
        if zip_files:
            col_zip1, col_zip2, col_zip3 = st.columns(3)
            with col_zip1:
                if st.button("Zip Worlds"):
                    create_zip_of_files([w['filename'] for w in get_saved_worlds()], "Worlds")
            with col_zip2:
                if st.button("Zip Chats"):
                    create_zip_of_files(glob.glob(os.path.join(CHAT_DIR, "*.md")), "Chats")
            with col_zip3:
                if st.button("Zip Audio"):
                    create_zip_of_files(glob.glob(os.path.join(AUDIO_DIR, "*.mp3")) + glob.glob(os.path.join(AUDIO_CACHE_DIR, "*.mp3")), "Audio")
            st.caption("Existing Zip Files:")
            for zip_file in zip_files:
                st.markdown(get_download_link(zip_file, "zip"), unsafe_allow_html=True)
        else:
            st.caption("No zip archives found.")

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

def initialize_world():
    if not st.session_state.get('initial_world_state_loaded', False):
        print("Performing initial world load for session...")
        saved_worlds = get_saved_worlds()
        loaded_successfully = False
        if saved_worlds:
            latest_world_file_basename = os.path.basename(saved_worlds[0]['filename'])
            print(f"Loading most recent world on startup: {latest_world_file_basename}")
            if load_world_state_from_md(latest_world_file_basename):
                loaded_successfully = True
            else:
                print("Failed to load most recent world, starting empty.")
        else:
            print("No saved worlds found, starting with empty state.")
        if not loaded_successfully:
            with world_objects_lock:
                world_objects.clear()
            st.session_state.current_world_file = None
        st.session_state.initial_world_state_loaded = True
        print("Initial world load process complete.")

if __name__ == "__main__":
    init_session_state()
    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
    initialize_world()
    render_sidebar()
    render_main_content()