app.py

# app.py (Merged Version)
import streamlit as st
import asyncio
import websockets
import uuid
from datetime import datetime
import os
import random
import time
import hashlib
# from PIL import Image # Keep commented unless needed for image pasting->3D texture?
import glob
import base64
import io
import streamlit.components.v1 as components
import edge_tts
# from audio_recorder_streamlit import audio_recorder # Keep commented unless re-adding audio input
import nest_asyncio
import re
import pytz
import shutil
# import anthropic # Keep commented unless integrating Claude
# import openai # Keep commented unless integrating OpenAI
from PyPDF2 import PdfReader
import threading
import json
import zipfile
# from gradio_client import Client # Keep commented unless integrating ArXiv/Gradio
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 # Still needed for some UI interactions

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

# 🎨 Page Config (From New App)
st.set_page_config(
    page_title="🤖🏗️ Shared World Builder 🏆",
    page_icon="🏗️",
    layout="wide",
    initial_sidebar_state="expanded" # Keep sidebar open initially
)

# --- Constants (Combined & 3D Added) ---
# Chat/User Constants
icons = '🤖🏗️🗣️' # Updated icons
Site_Name = '🤖🏗️ Shared World Builder 🗣️'
START_ROOM = "World Lobby 🌍"
FUN_USERNAMES = { # Simplified for clarity, can expand later
    "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()))
FILE_EMOJIS = {"md": "📝", "mp3": "🎵", "png": "🖼️", "mp4": "🎥", "zip": "📦", "csv":"📄"}

# 3D World Constants
SAVE_DIR = "saved_worlds"
PLOT_WIDTH = 50.0
PLOT_DEPTH = 50.0
CSV_COLUMNS = ['obj_id', 'type', 'pos_x', 'pos_y', 'pos_z', 'rot_x', 'rot_y', 'rot_z', 'rot_order']
WORLD_STATE_FILE = "world_state.json" # Using JSON for simpler in-memory<->disk state

# --- Directories (Combined) ---
for d in ["chat_logs", "audio_logs", "audio_cache", SAVE_DIR]: # Added SAVE_DIR
    os.makedirs(d, exist_ok=True)

CHAT_DIR = "chat_logs"
MEDIA_DIR = "." # Where general files are saved/served from
AUDIO_CACHE_DIR = "audio_cache"
AUDIO_DIR = "audio_logs"
STATE_FILE = "user_state.txt" # For remembering username

CHAT_FILE = os.path.join(CHAT_DIR, "global_chat.md")
# Removed vote files for simplicity

# --- API Keys (Keep placeholder logic) ---
load_dotenv()
# anthropic_key = os.getenv('ANTHROPIC_API_KEY', st.secrets.get('ANTHROPIC_API_KEY', ""))
# openai_api_key = os.getenv('OPENAI_API_KEY', st.secrets.get('OPENAI_API_KEY', ""))
# openai_client = openai.OpenAI(api_key=openai_api_key)

# --- Helper Functions (Combined & Adapted) ---

def format_timestamp_prefix(username=""):
    # Using UTC for consistency in logs/filenames across timezones potentially
    now = datetime.now(pytz.utc)
    # Simplified format
    return f"{now.strftime('%Y%m%d_%H%M%S')}_{username}"

# --- Performance Timer (Optional, Keep if desired) ---
class PerformanceTimer:
    # ... (keep class as is from new app.py if needed) ...
    pass

# --- 3D World State Management (Adapted from original + WebSocket focus) ---

# Global structure to hold the current state of the world IN MEMORY
# Use defaultdict for easier adding
# Needs thread safety if accessed by multiple websocket handlers simultaneously.
# For now, relying on Streamlit's single-thread-per-session execution
# and assuming broadcast updates are okay without strict locking for this scale.
# A lock would be needed for production robustness.
# world_objects_lock = threading.Lock() # Import threading if using lock
world_objects = defaultdict(dict) # Holds {obj_id: object_data}

def load_world_state_from_disk():
    """Loads world state from the JSON file or fallback to CSVs."""
    global world_objects
    loaded_count = 0
    if os.path.exists(WORLD_STATE_FILE):
        try:
            with open(WORLD_STATE_FILE, 'r') as f:
                data = json.load(f)
                # Ensure keys are strings if they got saved as ints somehow
                world_objects = defaultdict(dict, {str(k): v for k, v in data.items()})
                loaded_count = len(world_objects)
                print(f"Loaded {loaded_count} objects from {WORLD_STATE_FILE}")
        except json.JSONDecodeError:
            print(f"Error reading {WORLD_STATE_FILE}. Falling back to CSVs.")
            world_objects = defaultdict(dict) # Reset before loading from CSV
        except Exception as e:
            print(f"Error loading from {WORLD_STATE_FILE}: {e}. Falling back to CSVs.")
            world_objects = defaultdict(dict) # Reset

    # Fallback or initial load from CSVs if JSON fails or doesn't exist
    if not world_objects:
         print("Loading world state from CSV files...")
         # Use the cached CSV loading logic, but populate the global dict
         loaded_from_csv = get_all_world_objects_from_csv() # Gets list
         for obj in loaded_from_csv:
             world_objects[obj['obj_id']] = obj
         loaded_count = len(world_objects)
         print(f"Loaded {loaded_count} objects from CSVs.")
         # Save immediately to JSON for next time
         save_world_state_to_disk()

    return loaded_count

def save_world_state_to_disk():
    """Saves the current in-memory world state to a JSON file."""
    global world_objects
    print(f"Saving {len(world_objects)} objects to {WORLD_STATE_FILE}...")
    try:
        # with world_objects_lock: # Use lock if implementing thread safety
        with open(WORLD_STATE_FILE, 'w') as f:
             # Convert defaultdict back to regular dict for saving
             json.dump(dict(world_objects), f, indent=2)
        print("World state saved successfully.")
        return True
    except Exception as e:
        print(f"Error saving world state to {WORLD_STATE_FILE}: {e}")
        st.error(f"Failed to save world state: {e}")
        return False

# --- Functions to load from CSVs (kept for initial load/fallback) ---
@st.cache_data(ttl=3600)
def load_plot_metadata():
    """Scans save dir for plot_X*_Z*.csv, sorts, calculates metadata."""
    # ... (Keep function as is from original app.py) ...
    print(f"[{time.time():.2f}] Loading plot metadata...")
    plot_files = []
    try:
        plot_files = [f for f in os.listdir(SAVE_DIR) if f.endswith(".csv") and f.startswith("plot_X")]
    except FileNotFoundError: return []
    except Exception as e: return []

    parsed_plots = []
    for filename in plot_files:
        try:
            file_path = os.path.join(SAVE_DIR, filename)
            # Basic check for empty file before parsing name
            if not os.path.exists(file_path) or os.path.getsize(file_path) <= 2: continue

            parts = filename[:-4].split('_')
            grid_x = int(parts[1][1:])
            grid_z = int(parts[2][1:])
            plot_name = " ".join(parts[3:]) if len(parts) > 3 else f"Plot ({grid_x},{grid_z})"
            parsed_plots.append({
                'id': filename[:-4], 'filename': filename,
                'grid_x': grid_x, 'grid_z': grid_z, 'name': plot_name,
                'x_offset': grid_x * PLOT_WIDTH, 'z_offset': grid_z * PLOT_DEPTH
            })
        except Exception as e:
            st.warning(f"Error parsing metadata from filename '{filename}': {e}. Skipping.")
            continue
    parsed_plots.sort(key=lambda p: (p['grid_x'], p['grid_z']))
    return parsed_plots

def load_single_plot_objects_relative(filename):
    """Loads objects from a specific CSV file, keeping coordinates relative."""
    # ... (Keep function as is from previous merged version, including validation) ...
    file_path = os.path.join(SAVE_DIR, filename)
    try:
        if not os.path.exists(file_path) or os.path.getsize(file_path) == 0: return []
        df = pd.read_csv(file_path)
        if df.empty: return []
        if 'obj_id' not in df.columns: df['obj_id'] = [str(uuid.uuid4()) for _ in range(len(df))]
        else: df['obj_id'] = df['obj_id'].fillna(pd.Series([str(uuid.uuid4()) for _ in range(len(df))])).astype(str)
        for col in ['type', 'pos_x', 'pos_y', 'pos_z']:
            if col not in df.columns: return []
        for col, default in [('rot_x', 0.0), ('rot_y', 0.0), ('rot_z', 0.0), ('rot_order', 'XYZ')]:
            if col not in df.columns: df[col] = default
        df.fillna({'rot_x': 0.0, 'rot_y': 0.0, 'rot_z': 0.0, 'rot_order': 'XYZ'}, inplace=True)
        for col in ['pos_x', 'pos_y', 'pos_z', 'rot_x', 'rot_y', 'rot_z']:
             df[col] = pd.to_numeric(df[col], errors='coerce')
        df.dropna(subset=['pos_x', 'pos_y', 'pos_z'], inplace=True)
        df['type'] = df['type'].astype(str)
        return df[CSV_COLUMNS].to_dict('records')
    except Exception as e: return []


@st.cache_data(show_spinner="Loading initial world objects from CSVs...")
def get_all_world_objects_from_csv():
    """Loads ALL objects from ALL known plots into world coordinates FROM CSVs."""
    # ... (Keep function as is from previous merged version) ...
    print(f"[{time.time():.2f}] Reloading ALL world objects from CSV files...")
    all_objects = {}
    plots_meta = load_plot_metadata()
    for plot in plots_meta:
        relative_objects = load_single_plot_objects_relative(plot['filename'])
        for obj in relative_objects:
            obj_id = obj.get('obj_id')
            if not obj_id: continue
            world_obj = {
                'obj_id': obj_id, 'type': obj.get('type', 'Unknown'),
                'position': {'x': obj.get('pos_x', 0.0) + plot['x_offset'], 'y': obj.get('pos_y', 0.0), 'z': obj.get('pos_z', 0.0) + plot['z_offset']},
                'rotation': {'_x': obj.get('rot_x', 0.0), '_y': obj.get('rot_y', 0.0), '_z': obj.get('rot_z', 0.0), '_order': obj.get('rot_order', 'XYZ')}
            }
            all_objects[obj_id] = world_obj
    return list(all_objects.values())

# --- Session State Init (Combined & Expanded) ---
def init_session_state():
    defaults = {
        # From Chat App
        'server_running': False, 'server_task': None, 'active_connections': defaultdict(dict), # Use defaultdict
        'last_chat_update': 0, 'message_text': "", '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': True, 'last_message': "", 'timer_start': time.time(),
         'last_sent_transcript': "", 'last_refresh': time.time(),
         'auto_refresh': False, # Default auto-refresh off for world builder?
         'refresh_rate': 30, # Default refresh rate

        # From 3D World App (or adapted)
        'selected_object': 'None', # Current building tool
        # 'world_objects': defaultdict(dict), # In-memory state now global 'world_objects'
        'initial_world_state_loaded': False, # Flag to load state only once

        # Keep others if needed, removed some for clarity
        'operation_timings': {}, 'performance_metrics': defaultdict(list),
    }
    for k, v in defaults.items():
        if k not in st.session_state:
            st.session_state[k] = v
    # Ensure nested dicts are present
    if 'marquee_settings' not in st.session_state: st.session_state.marquee_settings = defaults['marquee_settings']
    if 'active_connections' not in st.session_state: st.session_state.active_connections = defaultdict(dict)

# --- Marquee Helpers (Keep from New App) ---
def update_marquee_settings_ui(): # ... (keep function as is) ...
    pass # Placeholder if not immediately needed
def display_marquee(text, settings, key_suffix=""): # ... (keep function as is) ...
    pass # Placeholder

# --- Text & File Helpers (Keep & Adapt from New App) ---
def clean_text_for_tts(text): # ... (keep function as is) ...
    return re.sub(r'[#*!\[\]]+', '', ' '.join(text.split()))[:200] or "No text"

def generate_filename(prompt, username, file_type="md", title=None): # ... (keep function as is) ...
     timestamp = format_timestamp_prefix(username)
     # Simplified filename generation
     base = clean_text_for_filename(title if title else prompt[:30])
     hash_val = hashlib.md5(prompt.encode()).hexdigest()[:6]
     return f"{timestamp}_{base}_{hash_val}.{file_type}"

def clean_text_for_filename(text): # ... (keep function as is) ...
    return '_'.join(re.sub(r'[^\w\s-]', '', text.lower()).split())[:50]

def create_file(content, username, file_type="md", title=None): # ... (keep function as is) ...
    filename = generate_filename(content, username, file_type, title)
    # Ensure saving to correct directory based on type? Assume current dir for now
    save_path = filename # os.path.join(MEDIA_DIR, filename)?
    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, file_type="mp3"): # ... (keep function as is, ensure FILE_EMOJIS updated) ...
    cache_key = f"dl_{file}_{os.path.getmtime(file) if os.path.exists(file) else 0}"
    if cache_key not in st.session_state['download_link_cache']:
        if not os.path.exists(file): return "File not found"
        with open(file, "rb") as f: b64 = base64.b64encode(f.read()).decode()
        mime_types = {"mp3": "audio/mpeg", "png": "image/png", "mp4": "video/mp4", "md": "text/markdown", "zip": "application/zip", "csv": "text/csv", "json": "application/json"}
        st.session_state['download_link_cache'][cache_key] = f'<a href="data:{mime_types.get(file_type, "application/octet-stream")};base64,{b64}" download="{os.path.basename(file)}">{FILE_EMOJIS.get(file_type, "📄")} Download {os.path.basename(file)}</a>'
    return st.session_state['download_link_cache'][cache_key]

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

def load_username(): # ... (keep function as is) ...
    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 to load username: {e}")
    return None

# --- Audio Processing (Keep from New App) ---
async def async_edge_tts_generate(text, voice, username): # Simplified args
    # ... (keep core logic, maybe save to AUDIO_DIR) ...
    cache_key = f"{text[:100]}_{voice}"
    if cache_key in st.session_state['audio_cache']: return st.session_state['audio_cache'][cache_key]
    text = clean_text_for_tts(text)
    if not text or text == "No text": return None
    filename_base = generate_filename(text, username, "mp3")
    save_path = os.path.join(AUDIO_DIR, filename_base)
    try:
        communicate = edge_tts.Communicate(text, 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: return None
    except Exception as e: return None

def play_and_download_audio(file_path): # ... (keep function as is) ...
    if file_path and os.path.exists(file_path):
        st.audio(file_path)
        file_type = file_path.split('.')[-1]
        st.markdown(get_download_link(file_path, file_type), unsafe_allow_html=True)

# --- Chat Saving/Loading (Keep & Adapt from New App) ---
async def save_chat_entry(username, message, voice, is_markdown=False):
    # ... (keep core logic, save to CHAT_DIR) ...
    if not message.strip(): return None, None
    central = pytz.timezone('US/Central') # Or use UTC
    timestamp = datetime.now(central).strftime("%Y-%m-%d %H:%M:%S")
    entry = f"[{timestamp}] {username} ({voice}): {message}" if not is_markdown else f"[{timestamp}] {username} ({voice}):\n```markdown\n{message}\n```"
    md_filename_base = generate_filename(message, username, "md")
    md_file = create_file(entry, username, "md", os.path.join(CHAT_DIR, md_filename_base)) # Save to chat_logs
    # Simplified - don't write to global CHAT_FILE on every message, maybe periodically?
    # Append to session state history for immediate display
    st.session_state.chat_history.append(entry)
    # Generate audio
    audio_file = None
    if st.session_state.get('enable_audio', True): # Check if enabled
        audio_file = await async_edge_tts_generate(message, voice, username)
    return md_file, audio_file

async def load_chat_history():
     # Load from individual files in CHAT_DIR for resilience? Or just session state?
     # For now, rely on session state + initial load from files if needed.
     if not st.session_state.chat_history:
          chat_files = sorted(glob.glob(os.path.join(CHAT_DIR, "*.md")), key=os.path.getmtime)
          for f in chat_files:
               try:
                    with open(f, 'r', encoding='utf-8') as file:
                         st.session_state.chat_history.append(file.read().strip())
               except Exception: pass # Ignore read errors
     return st.session_state.chat_history


# --- WebSocket Handling (Adapted for 3D State) ---
# Global set to track connected client IDs for efficient broadcast checks
connected_clients = set()

async def websocket_handler(websocket, path):
    client_id = str(websocket.id) # Use websocket's built-in ID
    connected_clients.add(client_id)
    username = st.session_state.get('username', f"User_{client_id[:4]}")
    print(f"Client connected: {client_id} ({username})")

    # Send initial world state to the new client
    try:
        # with world_objects_lock: # Lock if using threads
        initial_state_msg = json.dumps({
            "type": "initial_state",
            "payload": dict(world_objects) # Send current world state
        })
        await websocket.send(initial_state_msg)
        print(f"Sent initial state ({len(world_objects)} objects) to {client_id}")

        # Announce join (optional)
        await broadcast_message(json.dumps({
             "type": "user_join",
             "payload": {"username": username, "id": client_id}
        }), exclude_id=client_id) # Don't send to self

    except Exception as e:
        print(f"Error sending initial state to {client_id}: {e}")

    # Main message loop
    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) # Get username from payload or default

                if msg_type == "chat_message":
                    print(f"Received chat from {sender_username}: {payload.get('message')}")
                    voice = FUN_USERNAMES.get(sender_username, "en-US-AriaNeural")
                    # Save chat locally (optional async call)
                    asyncio.create_task(save_chat_entry(sender_username, payload.get('message', ''), voice))
                    # Broadcast chat to others
                    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:
                        print(f"Received place_object from {sender_username}: {obj_data.get('type')} ({obj_data['obj_id']})")
                        # with world_objects_lock: # Lock if needed
                        world_objects[obj_data['obj_id']] = obj_data # Add/update in memory
                        # Broadcast placement to others
                        broadcast_payload = json.dumps({
                             "type": "object_placed",
                             "payload": {"object_data": obj_data, "username": sender_username}
                        })
                        await broadcast_message(broadcast_payload, exclude_id=client_id)
                        # Maybe trigger periodic save here? Or rely on manual save.
                    else:
                         print(f"Invalid place_object payload from {client_id}")

                elif msg_type == "delete_object":
                     obj_id = payload.get("obj_id")
                     if obj_id:
                          print(f"Received delete_object from {sender_username}: {obj_id}")
                          # with world_objects_lock: # Lock if needed
                          if obj_id in world_objects:
                              del world_objects[obj_id]
                              # Broadcast deletion
                              broadcast_payload = json.dumps({
                                   "type": "object_deleted",
                                   "payload": {"obj_id": obj_id, "username": sender_username}
                              })
                              await broadcast_message(broadcast_payload, exclude_id=client_id)

                # Add handlers for other types (player_move, request_save, etc.)

            except json.JSONDecodeError:
                print(f"Received invalid JSON from {client_id}: {message}")
            except Exception as e:
                print(f"Error processing message from {client_id}: {e}")

    except websockets.ConnectionClosedOK:
        print(f"Client disconnected normally: {client_id} ({username})")
    except websockets.ConnectionClosedError as e:
        print(f"Client connection closed with error: {client_id} ({username}) - {e}")
    finally:
        connected_clients.discard(client_id)
        # Announce leave (optional)
        await broadcast_message(json.dumps({
             "type": "user_leave",
             "payload": {"username": username, "id": client_id}
        }))
        print(f"Client disconnected: {client_id} ({username}). Remaining: {len(connected_clients)}")


# Modified broadcast to use the global set and skip sender
async def broadcast_message(message, exclude_id=None):
    tasks = []
    disconnected_clients_this_call = set()

    # Iterate over a copy of the client set in case it changes during iteration
    current_client_ids = list(connected_clients)

    for client_id in current_client_ids:
        if client_id == exclude_id:
            continue

        websocket = st.session_state.active_connections.get(client_id) # Get WS object
        if websocket:
             try:
                  # Create task for sending; allows concurrent sends
                  tasks.append(asyncio.create_task(websocket.send(message)))
             except websockets.ConnectionClosed:
                  print(f"Found disconnected client during broadcast prep: {client_id}")
                  disconnected_clients_this_call.add(client_id)
             except RuntimeError as e: # Handles 'Event loop is closed' during shutdown
                  print(f"RuntimeError during broadcast prep for {client_id}: {e}")
                  disconnected_clients_this_call.add(client_id)
             except Exception as e:
                  print(f"Unexpected error during broadcast prep for {client_id}: {e}")
                  disconnected_clients_this_call.add(client_id)
        else:
            # Websocket object not found in session state, likely already disconnected
            disconnected_clients_this_call.add(client_id)

    # Wait for all send tasks to complete
    if tasks:
        await asyncio.gather(*tasks, return_exceptions=True) # Handle exceptions during send

    # Clean up disconnected clients found during this specific broadcast attempt
    if disconnected_clients_this_call:
         print(f"Cleaning up {len(disconnected_clients_this_call)} disconnected clients after broadcast.")
         for client_id in disconnected_clients_this_call:
              connected_clients.discard(client_id)
              st.session_state.active_connections.pop(client_id, None)


async def run_websocket_server():
    # Check if already running - basic flag protection
    if st.session_state.get('server_running_flag', False):
        print("Server already seems to be running or starting.")
        return
    st.session_state['server_running_flag'] = True
    print("Starting WebSocket server on 0.0.0.0:8765...")
    try:
        # Use localhost for testing, 0.0.0.0 for broader access (requires firewall config)
        server = await websockets.serve(websocket_handler, 'localhost', 8765)
        print("WebSocket server started successfully.")
        st.session_state['server_instance'] = server # Store server instance if needed for graceful shutdown
        await server.wait_closed() # Keep server running
    except OSError as e:
         print(f"### FAILED TO START WEBSOCKET SERVER: {e}")
         st.error(f"Failed to start WebSocket server on port 8765: {e}. Port might be in use.")
         # Try to connect as client if server fails? Or just stop.
    except Exception as e:
         print(f"### UNEXPECTED ERROR IN WEBSOCKET SERVER: {e}")
         st.error(f"An unexpected error occurred: {e}")
    finally:
         print("WebSocket server task finished.")
         st.session_state['server_running_flag'] = False
         st.session_state['server_instance'] = None


def start_websocket_server_thread():
    """Starts the WebSocket server in a separate thread."""
    # Check if thread already running
    if st.session_state.get('server_task') and st.session_state.server_task.is_alive():
        print("Server thread already running.")
        return
    print("Creating and starting new server thread.")
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    st.session_state.server_task = threading.Thread(target=loop.run_until_complete, args=(run_websocket_server(),), daemon=True)
    st.session_state.server_task.start()

# --- PDF to Audio (Keep if desired, maybe in a separate tab?) ---
class AudioProcessor: # ... (keep class as is) ...
      def __init__(self): self.cache_dir=AUDIO_CACHE_DIR; os.makedirs(self.cache_dir,exist_ok=True); self.metadata=json.load(open(f"{self.cache_dir}/metadata.json")) if os.path.exists(f"{self.cache_dir}/metadata.json") else {}
      def _save_metadata(self): #... (save logic) ...
            with open(f"{self.cache_dir}/metadata.json", 'w') as f: json.dump(self.metadata, f)
      async def create_audio(self, text, voice='en-US-AriaNeural'): # ... (audio creation logic) ...
            cache_key=hashlib.md5(f"{text}:{voice}".encode()).hexdigest(); cache_path=f"{self.cache_dir}/{cache_key}.mp3"
            if cache_key in self.metadata and os.path.exists(cache_path): return cache_path
            text=clean_text_for_tts(text); communicate=edge_tts.Communicate(text,voice); await communicate.save(cache_path)
            self.metadata[cache_key]={'timestamp': datetime.now().isoformat(), 'text_length': len(text), 'voice': voice}; self._save_metadata()
            return cache_path

def process_pdf(pdf_file, max_pages, voice, audio_processor): # ... (keep function as is) ...
      reader=PdfReader(pdf_file); total_pages=min(len(reader.pages),max_pages); texts,audios={}, {}
      async def process_page(i,text): audio_path=await audio_processor.create_audio(text,voice); audios[i]=audio_path
      for i in range(total_pages): text=reader.pages[i].extract_text(); texts[i]=text; threading.Thread(target=lambda: asyncio.run(process_page(i,text))).start()
      return texts,audios,total_pages

# --- ArXiv/AI Lookup (Commented out for focus) ---
# def parse_arxiv_refs(...): pass
# def generate_5min_feature_markdown(...): pass
# async def create_paper_audio_files(...): pass
# async def perform_ai_lookup(...): pass
# async def perform_claude_search(...): pass
# async def perform_arxiv_search(...): pass

# --- Voting (Removed for focus) ---
# def save_vote(...): pass
# def load_votes(...): pass

# --- Image Handling (Keep basic save, comment out Claude processing) ---
async def save_pasted_image(image, username): # Simplified
     img_hash = hashlib.md5(image.tobytes()).hexdigest()[:8]
     # Add check against existing hashes if needed: if img_hash in st.session_state.image_hashes: return None
     timestamp = format_timestamp_prefix(username)
     filename = f"{timestamp}_pasted_{img_hash}.png"
     filepath = os.path.join(MEDIA_DIR, filename) # Save in base dir
     try: image.save(filepath, "PNG"); return filepath
     except Exception as e: print(f"Failed image save: {e}"); return None

# --- Zip and Delete Files (Keep from New App) ---
def create_zip_of_files(files, prefix="Archive", query=""): # Simplified args
    if not files: return None
    timestamp = format_timestamp_prefix() # Generic timestamp
    zip_name = f"{prefix}_{timestamp}.zip"
    try:
        with zipfile.ZipFile(zip_name, 'w') as z:
            for f in files:
                 if os.path.exists(f): z.write(f, os.path.basename(f)) # Use basename in archive
        return zip_name
    except Exception as e: print(f"Zip creation failed: {e}"); return None

def delete_files(file_patterns, exclude_files=["README.md", STATE_FILE, WORLD_STATE_FILE]): # Takes list of patterns
    deleted_count = 0
    for pattern in file_patterns:
         # Be careful with glob patterns!
         files_to_delete = glob.glob(pattern)
         for f in files_to_delete:
             basename = os.path.basename(f)
             if basename not in exclude_files:
                  try: os.remove(f); deleted_count += 1
                  except Exception as e: print(f"Failed delete {f}: {e}")
    print(f"Deleted {deleted_count} files.")
    # Clear relevant caches?
    st.session_state['download_link_cache'] = {}


# --- Custom Paste Component (Keep from New App) ---
def paste_image_component(): # ... (Keep function as is) ...
      # Returns Image object, type string
      with st.form(key="paste_form"):
           paste_input = st.text_area("Paste Image Data Here (Ctrl+V)", key="paste_input_area", height=50)
           submit_button = st.form_submit_button("Paste Image 📋")
           if submit_button and paste_input and paste_input.startswith('data:image'):
               try:
                   mime_type = paste_input.split(';')[0].split(':')[1]; base64_str = paste_input.split(',')[1]
                   img_bytes = base64.b64decode(base64_str); img = Image.open(io.BytesIO(img_bytes))
                   st.image(img, caption=f"Pasted Image ({mime_type.split('/')[1].upper()})", width=150) # Smaller preview
                   return img, mime_type.split('/')[1]
               except Exception as e: st.error(f"Image decode error: {e}")
      return None, None


# --- Mapping Emojis to Primitive Types ---
# Ensure these types match the create[PrimitiveName] functions in index.html
PRIMITIVE_MAP = {
    "🌳": "Tree", "🗿": "Rock", "🏛️": "Simple House", "🌲": "Pine Tree", "🧱": "Brick Wall",
    "🔵": "Sphere", "📦": "Cube", " cylinder ": "Cylinder", "🍦": "Cone", "🍩": "Torus",
    "🍄": "Mushroom", "🌵": "Cactus", "🔥": "Campfire", "⭐": "Star", "💎": "Gem",
    "🗼": "Tower", "🚧": "Barrier", "⛲": "Fountain", "🏮": "Lantern", "푯": "Sign Post"
    # Add more pairs up to ~20
}

# --- Main Streamlit Interface ---
def main_interface():
    init_session_state()

    # --- Load initial world state ONCE per session ---
    if not st.session_state.initial_world_state_loaded:
        load_world_state_from_disk()
        st.session_state.initial_world_state_loaded = True

    # --- Username Setup ---
    saved_username = load_username()
    if saved_username and saved_username in FUN_USERNAMES:
        st.session_state.username = saved_username
    if not st.session_state.username:
        st.session_state.username = random.choice(list(FUN_USERNAMES.keys()))
        st.session_state.tts_voice = FUN_USERNAMES[st.session_state.username]
        save_username(st.session_state.username)
        # Don't automatically announce join here, let WebSocket handler do it on connect

    st.title(f"{Site_Name} - User: {st.session_state.username}")

    # --- Main Content Area ---
    tab_world, tab_chat, tab_files = st.tabs(["🏗️ World Builder", "🗣️ Chat", "📂 Files & Settings"])

    with tab_world:
        st.header("Shared 3D World")
        st.caption("Place objects using the sidebar tools. Changes are shared live!")

        # --- Embed HTML Component for Three.js ---
        html_file_path = 'index.html'
        try:
            with open(html_file_path, 'r', encoding='utf-8') as f:
                html_template = f.read()

            # Inject necessary data for JS: Username, WebSocket URL, initial state?
            # Initial state now sent via WebSocket, maybe don't inject here?
            # Let's inject username and WS url.
            ws_url = "ws://localhost:8765" # Use localhost for local dev

            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)}; // Send current tool
    window.PLOT_WIDTH = {json.dumps(PLOT_WIDTH)}; // Send constants needed by JS
    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}")

    with tab_chat:
        st.header(f"{START_ROOM} Chat")
        chat_history = asyncio.run(load_chat_history()) # Load history
        chat_container = st.container(height=500) # Scrollable chat area
        with chat_container:
             # Display chat history (most recent at bottom)
             for entry in reversed(chat_history[-50:]): # Show last 50 messages
                  st.markdown(entry) # Use markdown to render potential code blocks

        # Chat Input Area
        message = st.text_input("Your Message:", key="message_input", label_visibility="collapsed")
        if st.button("Send Chat 💬", key="send_chat_button") or (st.session_state.autosend and message):
             if message.strip() and message != st.session_state.last_message:
                 st.session_state.last_message = message
                 voice = FUN_USERNAMES.get(st.session_state.username, "en-US-AriaNeural")
                 # Send via WebSocket
                 ws_message = json.dumps({
                      "type": "chat_message",
                      "payload": {"username": st.session_state.username, "message": message, "voice": voice}
                 })
                 # Run broadcast in asyncio task to avoid blocking Streamlit
                 asyncio.run(broadcast_message(ws_message))
                 # Save locally (optional now, as broadcast handles real-time)
                 asyncio.run(save_chat_entry(st.session_state.username, message, voice))
                 # Clear input - needs rerun or JS callback
                 st.session_state.message_input = ""
                 st.rerun() # Force rerun to clear input and update display

    with tab_files:
        st.header("File Management & Settings")
        # Add options from the new app's sidebar here if desired
        # e.g., Zipping, Deleting, Marquee settings

        st.subheader("Server & World State")
        col_ws, col_save = st.columns(2)
        with col_ws:
             ws_status = "Running" if st.session_state.get('server_running_flag', False) else "Stopped"
             st.metric("WebSocket Server", ws_status)
             st.metric("Connected Clients", len(connected_clients))
        with col_save:
             if st.button("💾 Save World State to Disk", key="save_world_disk"):
                  if save_world_state_to_disk():
                       st.success("World state saved!")
                  else:
                       st.error("Failed to save world state.")

        # Add file deletion buttons if needed
        st.subheader("Delete Files")
        col_del1, col_del2, col_del3 = st.columns(3)
        with col_del1:
             if st.button("🗑️ Delete Chats (.md)", key="del_chat_md"):
                  delete_files([os.path.join(CHAT_DIR, "*.md")])
                  st.session_state.chat_history = [] # Clear session history too
                  st.rerun()
        with col_del2:
            if st.button("🗑️ Delete Audio (.mp3)", 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()
        # Add more deletion options as needed


    # --- Sidebar Controls ---
    with st.sidebar:
        st.header("🏗️ Build Tools")
        st.caption("Select an object to place.")

        # --- Emoji Buttons for Primitives ---
        cols = st.columns(5) # Adjust grid width as needed
        col_idx = 0
        for emoji, name in PRIMITIVE_MAP.items():
            # Use button click to set selected_object
            button_key = f"primitive_{name}"
            # Highlight selected button? Could use custom CSS or just rely on state.
            button_type = "primary" if st.session_state.selected_object == name else "secondary"
            if cols[col_idx % 5].button(emoji, key=button_key, help=name, type=button_type):
                st.session_state.selected_object = name
                # Update JS selection without full rerun if possible
                try:
                    js_update_selection = f"updateSelectedObjectType({json.dumps(name)});"
                    streamlit_js_eval(js_code=js_update_selection, key="update_tool_js")
                except Exception as e:
                    print(f"Could not push tool update to JS: {e}")
                # Force a rerun to update button styles immediately
                st.rerun()
            col_idx += 1

        # Button to clear selection
        if st.button("🚫 Clear Tool", key="clear_tool"):
             if st.session_state.selected_object != 'None':
                  st.session_state.selected_object = 'None'
                  try: # Update JS too
                      streamlit_js_eval(js_code=f"updateSelectedObjectType('None');", key="update_tool_js_none")
                  except Exception: pass
                  st.rerun() # Rerun to update UI

        st.markdown("---")
        st.header("🗣️ Voice & User")
        # Username/Voice Selection
        new_username = st.selectbox("Change Name/Voice", list(FUN_USERNAMES.keys()), index=list(FUN_USERNAMES.keys()).index(st.session_state.username), key="username_select")
        if new_username != st.session_state.username:
            # Announce name change via WebSocket?
            change_msg = json.dumps({
                 "type":"user_rename",
                 "payload": {"old_username": st.session_state.username, "new_username": new_username}
            })
            asyncio.run(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()

        # Enable/Disable Audio Toggle
        st.session_state['enable_audio'] = st.toggle("Enable TTS Audio", value=st.session_state.get('enable_audio', True))

        st.markdown("---")
        st.info("Status and file management in 'Files & Settings' tab.")


# --- Main Execution ---
if __name__ == "__main__":
    init_session_state()
    # Start WebSocket server in a thread IF it's not already running
    if not st.session_state.get('server_task') or not st.session_state.server_task.is_alive():
         start_websocket_server_thread()
         time.sleep(1) # Give server a moment to start

    main_interface()