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 # Keep import if used
from collections import defaultdict, Counter
import pandas as pd # Keep for potential fallback logic if needed
from streamlit_js_eval import streamlit_js_eval
from PIL import Image # Needed for paste_image_component

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

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

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

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

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

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

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

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

# --- Mapping Emojis to Primitive Types ---
# Ensure these types match the createPrimitiveMesh function keys in index.html
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"
    # Add more pairs up to ~20 if desired
}

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

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

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

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

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


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

def run_async(async_func, *args, **kwargs):
    """Runs an async function safely from a sync context using create_task."""
    try:
        loop = asyncio.get_running_loop()
        return loop.create_task(async_func(*args, **kwargs))
    except RuntimeError: # No running loop in this thread
        print(f"Warning: Running async func {async_func.__name__} in new event loop.")
        try:
            return asyncio.run(async_func(*args, **kwargs))
        except Exception as e:
            print(f"Error running async func {async_func.__name__} in new loop: {e}")
            return None
    except Exception as e:
        print(f"Error scheduling async task {async_func.__name__}: {e}")
        return None

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

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

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

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

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

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


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

```json
{json.dumps(world_data_dict, indent=2)}
```"""
        try:
            ensure_dir(SAVED_WORLDS_DIR)
            with open(save_path, 'w', encoding='utf-8') as f: f.write(md_content)
            print(f"World state saved successfully to {target_filename_base}")
            success = True
        except Exception as e:
            print(f"Error saving world state to {save_path}: {e}")
    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\S]*?\})\s*```", content, re.IGNORECASE) # More robust regex
        if not json_match: st.error(f"Could not find valid JSON block in {filename_base}"); return False

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

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

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

def get_saved_worlds():
    """Scans the saved worlds directory for world MD files and parses them."""
    try:
        ensure_dir(SAVED_WORLDS_DIR)
        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]
        parsed_worlds.sort(key=lambda x: x['dt'] if x['dt'] else datetime.min.replace(tzinfo=pytz.utc), reverse=True)
        return parsed_worlds
    except Exception as e:
         print(f"Error scanning saved worlds: {e}"); st.error(f"Could not scan saved worlds: {e}"); return []

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

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

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

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

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

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

def create_file(content, username, file_type="md", save_path=None):
    if not save_path: filename = generate_filename(content, username, file_type); save_path = os.path.join(MEDIA_DIR, filename)
    ensure_dir(os.path.dirname(save_path))
    try:
        with open(save_path, 'w', encoding='utf-8') as f: f.write(content)
        # print(f"Created file: {save_path}"); # Can be too verbose
        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>")

# --- Audio / TTS ---
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}")

# --- Chat ---
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) # Save to file
    if 'chat_history' not in st.session_state: st.session_state.chat_history = [];
    st.session_state.chat_history.append(entry) # Add to live history
    audio_file = None;
    if st.session_state.get('enable_audio', True):
         tts_message = message # Use original message for TTS
         audio_file = await async_edge_tts_generate(tts_message, voice, username)
    return md_file, audio_file

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

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

def delete_files(file_patterns, exclude_files=None):
    """Deletes files matching patterns, excluding protected/specified files."""
    protected = [STATE_FILE, "app.py", "index.html", "requirements.txt", "README.md"]
    # Don't automatically protect current world file during generic delete
    if exclude_files: protected.extend(exclude_files)

    deleted_count = 0; errors = 0
    for pattern in file_patterns:
         pattern_path = pattern # Assume pattern includes path from os.path.join
         print(f"Attempting to delete files matching: {pattern_path}")
         try:
              files_to_delete = glob.glob(pattern_path)
              if not files_to_delete: print(f"No files found for pattern: {pattern}"); continue
              for f_path in files_to_delete:
                   basename = os.path.basename(f_path)
                   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)
    elif deleted_count > 0: st.success(msg)
    else: st.info("No matching files found to delete.")
    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
      # Added default value to text_area to ensure key exists even if empty
      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"): # Simplified trigger
           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 image if processed successfully in this run
      return pasted_img


# --- PDF Processing ---
class AudioProcessor:
      def __init__(self): self.cache_dir=AUDIO_CACHE_DIR; ensure_dir(self.cache_dir); self.metadata=json.load(open(f"{self.cache_dir}/metadata.json", 'r')) if os.path.exists(f"{self.cache_dir}/metadata.json") else {}
      def _save_metadata(self):
            try:
                 with open(f"{self.cache_dir}/metadata.json", 'w') as f: json.dump(self.metadata, f, indent=2)
            except Exception as e: print(f"Failed metadata save: {e}")
      async def create_audio(self, text, voice='en-US-AriaNeural'):
            cache_key=hashlib.md5(f"{text[:150]}:{voice}".encode()).hexdigest(); cache_path=os.path.join(self.cache_dir, f"{cache_key}.mp3");
            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") # Change header
    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:
                # asyncio.run can cause issues if called repeatedly in threads, use run_async helper
                audio_path_future = run_async(run_async_audio)
                # This needs careful handling. run_async might return a task or result.
                # For simplicity here, let's assume it blocks or we retrieve result later.
                # A better pattern uses concurrent.futures or manages event loop properly.
                # Sticking to asyncio.run for now as run_async tries it.
                audio_path = asyncio.run(run_async_audio()) # Revert to simpler asyncio.run for thread context
                if audio_path:
                     with results_lock: audios[page_num] = audio_path
            except Exception as page_e: print(f"Err process page {page_num+1}: {page_e}")

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

        # Wait for threads and display progress
        progress_bar = st.progress(0.0, 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 # 10 min timeout
            time.sleep(0.5)
        progress_bar.progress(1.0, text="Processing complete.")

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

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

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

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

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

async def send_safely(websocket, message, client_id):
     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: # Send initial state
        with world_objects_lock: initial_state_payload = dict(world_objects)
        initial_state_msg = json.dumps({"type": "initial_state", "payload": initial_state_payload}); await websocket.send(initial_state_msg)
        print(f"Sent initial state ({len(initial_state_payload)} objs) to {client_id}")
        await broadcast_message(json.dumps({"type": "user_join", "payload": {"username": username, "id": client_id}}), exclude_id=client_id)
    except Exception as e: print(f"Error initial phase {client_id}: {e}")

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

                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)
                    else: print(f"WS Invalid place_object payload: {payload}")

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

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

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


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():
    """Renders the Streamlit sidebar contents."""
    with st.sidebar:
        st.header("💾 World Versions")
        st.caption("Load or save named world states.")

        saved_worlds = get_saved_worlds()
        world_options_display = {os.path.basename(w['filename']): f"{w['name']} ({w['timestamp']})" 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: # Show expander after limit
                 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 # Stop outer loop after showing expander

            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.")
        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 name_arg=name: streamlit_js_eval(f"updateSelectedObjectType({json.dumps(name_arg)});", key=f"update_tool_js_{name_arg}"))
                    st.rerun()
            col_idx += 1
        st.markdown("---")
        if st.button("🚫 Clear Tool", key="clear_tool", use_container_width=True):
             if st.session_state.get('selected_object', 'None') != 'None':
                  st.session_state.selected_object = 'None';
                  run_async(lambda: streamlit_js_eval("updateSelectedObjectType('None');", key="update_tool_js_none"))
                  st.rerun()

        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_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)")

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

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

    # --- Chat Tab ---
    with tab_chat:
        st.header(f"{START_ROOM} Chat")
        chat_history_task = run_async(load_chat_history) # Schedule load
        chat_history = chat_history_task.result() if chat_history_task else st.session_state.chat_history # Get result if task ran sync, else use current state
        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 = "" # 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")

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

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

        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:
             # Ensure correct indentation here for the else block
             st.caption("No zip archives found.")


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

def initialize_world():
    """Loads initial world state (most recent) if not already done for this session."""
    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__":
    # 1. Initialize session state
    init_session_state()

    # 2. Start WebSocket server thread if needed
    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

    # 3. Load initial world state (once per session)
    initialize_world()

    # 4. Render UI
    render_sidebar()
    render_main_content()