import streamlit as st
import asyncio
import websockets
import uuid
from datetime import datetime
import os
import random
import hashlib
import base64
import edge_tts
import nest_asyncio
import re
import threading
import time
import json
import streamlit.components.v1 as components
from gradio_client import Client
from streamlit_marquee import streamlit_marquee
import folium
from streamlit_folium import folium_static

# Patch asyncio for nesting
nest_asyncio.apply()

# Page Config
st.set_page_config(
    layout="wide",
    page_title="Rocky Mountain Quest 3D 🏔️🎮",
    page_icon="🦌"
)

# Game Config
GAME_NAME = "Rocky Mountain Quest 3D 🏔️🎮"
START_LOCATION = "Trailhead Camp ⛺"
CHARACTERS = [
    {"name": "Trailblazer Tim 🌄", "voice": "en-US-GuyNeural", "desc": "Fearless hiker seeking epic trails!", "color": 0x00ff00},
    {"name": "Meme Queen Mia 😂", "voice": "en-US-JennyNeural", "desc": "Spreads laughs with wild memes!", "color": 0xff00ff},
    {"name": "Elk Whisperer Eve 🦌", "voice": "en-GB-SoniaNeural", "desc": "Talks to wildlife, loves nature!", "color": 0x0000ff},
    {"name": "Tech Titan Tara 💾", "voice": "en-AU-NatashaNeural", "desc": "Codes her way through the Rockies!", "color": 0xffff00},
    {"name": "Ski Guru Sam ⛷️", "voice": "en-CA-ClaraNeural", "desc": "Shreds slopes, lives for snow!", "color": 0xffa500},
    {"name": "Cosmic Camper Cal 🌠", "voice": "en-US-AriaNeural", "desc": "Stargazes and tells epic tales!", "color": 0x800080},
    {"name": "Rasta Ranger Rick 🍃", "voice": "en-GB-RyanNeural", "desc": "Chills with nature’s vibes!", "color": 0x00ffff},
    {"name": "Boulder Bro Ben 🪨", "voice": "en-AU-WilliamNeural", "desc": "Climbs rocks, bro-style!", "color": 0xff4500}
]
FILE_EMOJIS = {"md": "📜", "mp3": "🎵"}

# Prairie Simulator Locations
PRAIRIE_LOCATIONS = {
    "Deadwood, SD": (44.3769, -103.7298),
    "Wind Cave National Park": (43.6047, -103.4798),
    "Wyoming Spring Creek": (41.6666, -106.6666)
}

# Directories and Files
for d in ["chat_logs", "audio_logs"]:
    os.makedirs(d, exist_ok=True)

CHAT_DIR = "chat_logs"
AUDIO_DIR = "audio_logs"
STATE_FILE = "user_state.txt"
CHAT_FILE = os.path.join(CHAT_DIR, "quest_log.md")
GAME_STATE_FILE = "game_state.json"

# Cached Game State as "ML Model" using st.cache_resource
@st.cache_resource
def load_game_state(_timestamp):
    """Load or initialize the game state, treated as a cached resource."""
    if os.path.exists(GAME_STATE_FILE):
        with open(GAME_STATE_FILE, 'r') as f:
            state = json.load(f)
    else:
        state = {
            "players": {},
            "treasures": [
                {"id": str(uuid.uuid4()), "x": random.uniform(-40, 40), "z": random.uniform(-40, 40)},
                {"id": str(uuid.uuid4()), "x": random.uniform(-40, 40), "z": random.uniform(-40, 40)},
                {"id": str(uuid.uuid4()), "x": random.uniform(-40, 40), "z": random.uniform(-40, 40)},
                {"id": str(uuid.uuid4()), "x": random.uniform(-40, 40), "z": random.uniform(-40, 40)},
                {"id": str(uuid.uuid4()), "x": random.uniform(-40, 40), "z": random.uniform(-40, 40)}
            ],
            "history": []
        }
        with open(GAME_STATE_FILE, 'w') as f:
            json.dump(state, f)
    return state

def update_game_state(state):
    """Update the game state and persist to file."""
    with open(GAME_STATE_FILE, 'w') as f:
        json.dump(state, f)
    load_game_state.clear()  # Clear cache to force reload
    return load_game_state(time.time())

def reset_game_state():
    """Reset the game state to initial conditions."""
    state = {
        "players": {},
        "treasures": [
            {"id": str(uuid.uuid4()), "x": random.uniform(-40, 40), "z": random.uniform(-40, 40)},
            {"id": str(uuid.uuid4()), "x": random.uniform(-40, 40), "z": random.uniform(-40, 40)},
            {"id": str(uuid.uuid4()), "x": random.uniform(-40, 40), "z": random.uniform(-40, 40)},
            {"id": str(uuid.uuid4()), "x": random.uniform(-40, 40), "z": random.uniform(-40, 40)},
            {"id": str(uuid.uuid4()), "x": random.uniform(-40, 40), "z": random.uniform(-40, 40)}
        ],
        "history": []
    }
    return update_game_state(state)

# Helpers
def format_timestamp(username=""):
    now = datetime.now().strftime("%Y%m%d_%H%M%S")
    return f"{now}-by-{username}"

def clean_text_for_tts(text):
    return re.sub(r'[#*!\[\]]+', '', ' '.join(text.split()))[:200] or "No text"

def generate_filename(prompt, username, file_type="md"):
    timestamp = format_timestamp(username)
    hash_val = hashlib.md5(prompt.encode()).hexdigest()[:8]
    return f"{timestamp}-{hash_val}.{file_type}"

def create_file(prompt, username, file_type="md"):
    filename = generate_filename(prompt, username, file_type)
    with open(filename, 'w', encoding='utf-8') as f:
        f.write(prompt)
    return filename

def get_download_link(file, file_type="mp3"):
    with open(file, "rb") as f:
        b64 = base64.b64encode(f.read()).decode()
    mime_types = {"mp3": "audio/mpeg", "md": "text/markdown"}
    return f'<a href="data:{mime_types.get(file_type, "application/octet-stream")};base64,{b64}" download="{os.path.basename(file)}">{FILE_EMOJIS.get(file_type, "📥")} {os.path.basename(file)}</a>'

def save_username(username):
    with open(STATE_FILE, 'w') as f:
        f.write(username)

def load_username():
    if os.path.exists(STATE_FILE):
        with open(STATE_FILE, 'r') as f:
            return f.read().strip()
    return None

# Audio Processing
async def async_edge_tts_generate(text, voice, username):
    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)
    filename = f"{AUDIO_DIR}/{format_timestamp(username)}-{hashlib.md5(text.encode()).hexdigest()[:8]}.mp3"
    communicate = edge_tts.Communicate(text, voice)
    await communicate.save(filename)
    if os.path.exists(filename) and os.path.getsize(filename) > 0:
        st.session_state['audio_cache'][cache_key] = filename
        return filename
    return None

def play_and_download_audio(file_path):
    if file_path and os.path.exists(file_path):
        st.audio(file_path)
        st.markdown(get_download_link(file_path), unsafe_allow_html=True)

# WebSocket Broadcast
async def broadcast_message(message, room_id):
    if room_id in st.session_state.active_connections:
        disconnected = []
        for client_id, ws in st.session_state.active_connections[room_id].items():
            try:
                await ws.send(message)
            except websockets.ConnectionClosed:
                disconnected.append(client_id)
        for client_id in disconnected:
            if client_id in st.session_state.active_connections[room_id]:
                del st.session_state.active_connections[room_id][client_id]

# Chat and Quest Log
async def save_chat_entry(username, message, voice, is_markdown=False):
    if not message.strip() or message == st.session_state.get('last_transcript', ''):
        return None, None
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    entry = f"[{timestamp}] {username}: {message}" if not is_markdown else f"[{timestamp}] {username}:\n```markdown\n{message}\n```"
    md_file = create_file(entry, username, "md")
    with open(CHAT_FILE, 'a') as f:
        f.write(f"{entry}\n")
    audio_file = await async_edge_tts_generate(message, voice, username)
    await broadcast_message(f"{username}|{message}", "quest")
    st.session_state.chat_history.append(entry)
    st.session_state.last_transcript = message
    game_state = load_game_state(st.session_state.game_state_timestamp)
    game_state["history"].append(entry)
    if username in game_state["players"]:
        game_state["players"][username]["score"] += 10
        game_state["players"][username]["treasures"] += 1
        game_state["players"][username]["last_active"] = time.time()
    update_game_state(game_state)
    return md_file, audio_file

async def load_chat():
    if not os.path.exists(CHAT_FILE):
        with open(CHAT_FILE, 'a') as f:
            f.write(f"# {GAME_NAME} Log\n\nThe adventure begins at {START_LOCATION}! 🏔️\n")
    with open(CHAT_FILE, 'r') as f:
        content = f.read().strip()
    return content.split('\n')

# Session State Init
def init_session_state():
    defaults = {
        'server_running': False, 'server_task': None, 'active_connections': {},
        'chat_history': [], 'audio_cache': {}, 'last_transcript': "",
        'username': None, 'score': 0, 'treasures': 0, 'location': START_LOCATION,
        'speech_processed': False, 'players': {}, 'last_update': time.time(),
        'update_interval': 1, 'x_pos': 0, 'z_pos': 0, 'move_left': False,
        'move_right': False, 'move_up': False, 'move_down': False,
        'prairie_players': {}, 'last_chat_update': 0,
        'game_state_timestamp': time.time(), 'name_index': 0,
        'timeout': 60, 'auto_refresh': 30, 'last_activity': time.time()
    }
    for k, v in defaults.items():
        if k not in st.session_state:
            st.session_state[k] = v
    if st.session_state.username is None:
        st.session_state.name_index = (st.session_state.name_index + 1) % len(CHARACTERS)
        st.session_state.username = CHARACTERS[st.session_state.name_index]["name"]
        asyncio.run(save_chat_entry(st.session_state.username, "🗺️ Joins the Rocky Mountain Quest!", CHARACTERS[st.session_state.name_index]["voice"]))
        save_username(st.session_state.username)
    game_state = load_game_state(st.session_state.game_state_timestamp)
    if st.session_state.username not in game_state["players"]:
        game_state["players"][st.session_state.username] = {
            "x": random.uniform(-20, 20),
            "z": random.uniform(-40, 40),
            "color": CHARACTERS[st.session_state.name_index]["color"],
            "score": 0,
            "treasures": 0,
            "last_active": time.time()
        }
        update_game_state(game_state)

init_session_state()

# WebSocket for Multiplayer
async def websocket_handler(websocket, path):
    client_id = str(uuid.uuid4())
    room_id = "quest"
    if room_id not in st.session_state.active_connections:
        st.session_state.active_connections[room_id] = {}
    st.session_state.active_connections[room_id][client_id] = websocket
    username = st.session_state.username
    
    game_state = load_game_state(st.session_state.game_state_timestamp)
    if "prairie" in path:
        st.session_state.prairie_players[client_id] = {
            "username": username,
            "animal": "prairie_dog",
            "location": PRAIRIE_LOCATIONS["Deadwood, SD"],
            "color": CHARACTERS[st.session_state.name_index]["color"]
        }
        await broadcast_message(f"System|{username} joins the prairie!", room_id)
    else:
        if username not in game_state["players"]:
            game_state["players"][username] = {
                "x": random.uniform(-20, 20),
                "z": random.uniform(-40, 40),
                "color": CHARACTERS[st.session_state.name_index]["color"],
                "score": 0,
                "treasures": 0,
                "last_active": time.time()
            }
            update_game_state(game_state)
        st.session_state.players[client_id] = game_state["players"][username]
        await broadcast_message(f"System|{username} joins the quest!", room_id)
    
    try:
        async for message in websocket:
            if '|' in message:
                sender, content = message.split('|', 1)
                voice = next(c["voice"] for c in CHARACTERS if c["name"] == sender)
                game_state = load_game_state(st.session_state.game_state_timestamp)
                if content.startswith("MOVE:"):
                    _, x, z = content.split(":")
                    x, z = float(x), float(z)
                    st.session_state.players[client_id]["x"] = x
                    st.session_state.players[client_id]["z"] = z
                    game_state["players"][sender]["x"] = x
                    game_state["players"][sender]["z"] = z
                    game_state["players"][sender]["last_active"] = time.time()
                    update_game_state(game_state)
                elif content.startswith("SCORE:"):
                    score = int(content.split(":")[1])
                    st.session_state.players[client_id]["score"] = score
                    game_state["players"][sender]["score"] = score
                    game_state["players"][sender]["last_active"] = time.time()
                    update_game_state(game_state)
                elif content.startswith("TREASURE:"):
                    treasures = int(content.split(":")[1])
                    st.session_state.players[client_id]["treasures"] = treasures
                    game_state["players"][sender]["treasures"] = treasures
                    game_state["players"][sender]["last_active"] = time.time()
                    for i, t in enumerate(game_state["treasures"]):
                        if abs(st.session_state.players[client_id]["x"] - t["x"]) < 2 and \
                           abs(st.session_state.players[client_id]["z"] - t["z"]) < 2:
                            game_state["treasures"].pop(i)
                            break
                    update_game_state(game_state)
                elif content.startswith("PRAIRIE:"):
                    action, value = content.split(":", 1)
                    if action == "ANIMAL":
                        st.session_state.prairie_players[client_id]["animal"] = value
                    elif action == "MOVE":
                        target = PRAIRIE_LOCATIONS.get(value, PRAIRIE_LOCATIONS["Deadwood, SD"])
                        st.session_state.prairie_players[client_id]["location"] = target
                        action_msg = f"{sender} ({st.session_state.prairie_players[client_id]['animal']}) moves to {value}"
                        await save_chat_entry(sender, action_msg, voice)
                else:
                    await save_chat_entry(sender, content, voice)
                    await perform_arxiv_search(content, sender)
    except websockets.ConnectionClosed:
        await broadcast_message(f"System|{username} leaves the quest!", room_id)
        if client_id in st.session_state.players:
            del st.session_state.players[client_id]
        if client_id in st.session_state.prairie_players:
            del st.session_state.prairie_players[client_id]
    finally:
        if room_id in st.session_state.active_connections and client_id in st.session_state.active_connections[room_id]:
            del st.session_state.active_connections[room_id][client_id]

async def periodic_update():
    while True:
        if st.session_state.active_connections.get("quest"):
            game_state = load_game_state(st.session_state.game_state_timestamp)
            # Remove inactive players (timeout after 60 seconds)
            current_time = time.time()
            inactive_players = [p for p, data in game_state["players"].items() if current_time - data["last_active"] > st.session_state.timeout]
            for player in inactive_players:
                del game_state["players"][player]
                await broadcast_message(f"System|{player} timed out after 60 seconds!", "quest")
            update_game_state(game_state)

            player_list = ", ".join([p for p in game_state["players"].keys()]) or "No adventurers yet!"
            message = f"📢 Quest Update: Active Adventurers - {player_list}"
            player_data = json.dumps(list(game_state["players"].values()))
            await broadcast_message(f"MAP_UPDATE:{player_data}", "quest")
            
            prairie_list = ", ".join([f"{p['username']} ({p['animal']})" for p in st.session_state.prairie_players.values()]) or "No animals yet!"
            prairie_message = f"🌾 Prairie Update: {prairie_list}"
            prairie_data = json.dumps(list(st.session_state.prairie_players.values()))
            await broadcast_message(f"PRAIRIE_UPDATE:{prairie_data}", "quest")
            
            chat_content = await load_chat()
            await broadcast_message(f"CHAT_UPDATE:{json.dumps(chat_content[-10:])}", "quest")
            await broadcast_message(f"GAME_STATE:{json.dumps(game_state)}", "quest")
            await save_chat_entry("System", f"{message}\n{prairie_message}", "en-US-AriaNeural")
        await asyncio.sleep(st.session_state.update_interval)

async def run_websocket_server():
    if not st.session_state.get('server_running', False):
        server = await websockets.serve(websocket_handler, '0.0.0.0', 8765)
        st.session_state['server_running'] = True
        asyncio.create_task(periodic_update())
        await server.wait_closed()

def start_websocket_server():
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(run_websocket_server())

# ArXiv Integration
async def perform_arxiv_search(query, username):
    gradio_client = Client("awacke1/Arxiv-Paper-Search-And-QA-RAG-Pattern")
    refs = gradio_client.predict(
        query, 5, "Semantic Search", "mistralai/Mixtral-8x7B-Instruct-v0.1", api_name="/update_with_rag_md"
    )[0]
    result = f"📚 Ancient Rocky Knowledge:\n{refs}"
    voice = next(c["voice"] for c in CHARACTERS if c["name"] == username)
    md_file, audio_file = await save_chat_entry(username, result, voice, True)
    return md_file, audio_file

# Enhanced 3D Game HTML
rocky_map_html = f"""
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Rocky Mountain Quest 3D</title>
    <style>
        body {{ margin: 0; overflow: hidden; font-family: Arial, sans-serif; background: #000; }}
        #gameContainer {{ width: 800px; height: 600px; position: relative; }}
        canvas {{ width: 100%; height: 100%; display: block; }}
        #ui {{ 
            position: absolute; top: 10px; left: 10px; color: white; 
            background: rgba(0, 0, 0, 0.5); padding: 5px; border-radius: 5px; 
        }}
        #chatBox {{ 
            position: absolute; bottom: 60px; left: 10px; width: 300px; height: 150px; 
            background: rgba(0, 0, 0, 0.7); color: white; padding: 10px; 
            border-radius: 5px; overflow-y: auto; 
        }}
        #controls {{ 
            position: absolute; bottom: 10px; left: 10px; color: white; 
            background: rgba(0, 0, 0, 0.5); padding: 5px; border-radius: 5px; 
        }}
    </style>
</head>
<body>
    <div id="gameContainer">
        <div id="ui">
            <div id="players">Players: 1</div>
            <div id="score">Score: 0</div>
            <div id="treasures">Treasures: 0</div>
            <div id="timeout">Timeout: 60s</div>
        </div>
        <div id="chatBox"></div>
        <div id="controls">WASD/Arrows to move, Space to collect treasure</div>
    </div>

    <script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script>
    <script>
        const playerName = "{st.session_state.username}";
        let ws = new WebSocket('ws://localhost:8765');
        const scene = new THREE.Scene();
        const camera = new THREE.PerspectiveCamera(75, 800 / 600, 0.1, 1000);
        camera.position.set(0, 50, 50);
        camera.lookAt(0, 0, 0);

        const renderer = new THREE.WebGLRenderer({{ antialias: true }});
        renderer.setSize(800, 600);
        document.getElementById('gameContainer').appendChild(renderer.domElement);

        const ambientLight = new THREE.AmbientLight(0xffffff, 0.5);
        scene.add(ambientLight);
        const sunLight = new THREE.DirectionalLight(0xffddaa, 1);
        sunLight.position.set(50, 50, 50);
        sunLight.castShadow = true;
        scene.add(sunLight);

        const groundGeometry = new THREE.PlaneGeometry(100, 100);
        const groundMaterial = new THREE.MeshStandardMaterial({{ color: 0x228B22 }});
        const ground = new THREE.Mesh(groundGeometry, groundMaterial);
        ground.rotation.x = -Math.PI / 2;
        ground.receiveShadow = true;
        scene.add(ground);

        let players = {{}};
        const playerMeshes = {{}};
        let treasures = [];
        const treasureMeshes = {{}};
        let xPos = 0, zPos = 0, moveLeft = false, moveRight = false, moveUp = false, moveDown = false, collect = false;
        let score = 0, treasureCount = 0, lastActive = performance.now() / 1000;

        // Player initialization
        const playerGeometry = new THREE.BoxGeometry(2, 2, 2);
        const playerMaterial = new THREE.MeshPhongMaterial({{ color: {next(c["color"] for c in CHARACTERS if c["name"] == st.session_state.username)} }});
        const playerMesh = new THREE.Mesh(playerGeometry, playerMaterial);
        playerMesh.position.set(xPos, 1, zPos);
        playerMesh.castShadow = true;
        scene.add(playerMesh);
        players[playerName] = {{ mesh: playerMesh, score: 0, treasures: 0 }};

        // Treasure updating
        function updateTreasures(treasureData) {{
            treasureData.forEach(t => {{
                if (!treasureMeshes[t.id]) {{
                    const treasure = new THREE.Mesh(
                        new THREE.SphereGeometry(1, 8, 8),
                        new THREE.MeshPhongMaterial({{ color: 0xffff00 }})
                    );
                    treasure.position.set(t.x, 1, t.z);
                    treasure.castShadow = true;
                    treasureMeshes[t.id] = treasure;
                    treasures.push(treasure);
                    scene.add(treasure);
                }} else {{
                    treasureMeshes[t.id].position.set(t.x, 1, t.z);
                }}
            }});
            Object.keys(treasureMeshes).forEach(id => {{
                if (!treasureData.some(t => t.id === id)) {{
                    scene.remove(treasureMeshes[id]);
                    treasures = treasures.filter(t => t !== treasureMeshes[id]);
                    delete treasureMeshes[id];
                }}
            }});
        }}

        // Controls
        document.addEventListener('keydown', (event) => {{
            switch (event.code) {{
                case 'ArrowLeft': case 'KeyA': moveLeft = true; break;
                case 'ArrowRight': case 'KeyD': moveRight = true; break;
                case 'ArrowUp': case 'KeyW': moveUp = true; break;
                case 'ArrowDown': case 'KeyS': moveDown = true; break;
                case 'Space': collect = true; break;
            }}
            lastActive = performance.now() / 1000;
        }});
        document.addEventListener('keyup', (event) => {{
            switch (event.code) {{
                case 'ArrowLeft': case 'KeyA': moveLeft = false; break;
                case 'ArrowRight': case 'KeyD': moveRight = false; break;
                case 'ArrowUp': case 'KeyW': moveUp = false; break;
                case 'ArrowDown': case 'KeyS': moveDown = false; break;
                case 'Space': collect = false; break;
            }}
        }});

        function updatePlayer(delta) {{
            const speed = 20;
            if (moveLeft && xPos > -40) xPos -= speed * delta;
            if (moveRight && xPos < 40) xPos += speed * delta;
            if (moveUp && zPos > -40) zPos -= speed * delta;
            if (moveDown && zPos < 40) zPos += speed * delta;
            playerMesh.position.set(xPos, 1, zPos);
            ws.send(`${{playerName}}|MOVE:${{xPos}}:${{zPos}}`);

            if (collect) {{
                for (let i = treasures.length - 1; i >= 0; i--) {{
                    if (playerMesh.position.distanceTo(treasures[i].position) < 2) {{
                        const id = Object.keys(treasureMeshes).find(key => treasureMeshes[key] === treasures[i]);
                        scene.remove(treasures[i]);
                        treasures.splice(i, 1);
                        delete treasureMeshes[id];
                        score += 50;
                        treasureCount += 1;
                        ws.send(`${{playerName}}|SCORE:${{score}}`);
                        ws.send(`${{playerName}}|TREASURE:${{treasureCount}}`);
                    }}
                }}
            }}

            camera.position.set(xPos, 50, zPos + 50);
            camera.lookAt(xPos, 0, zPos);
            const timeout = 60 - (performance.now() / 1000 - lastActive);
            document.getElementById('timeout').textContent = `Timeout: ${{Math.max(0, timeout.toFixed(0))}}s`;
        }}

        function updatePlayers(playerData) {{
            playerData.forEach(player => {{
                if (!playerMeshes[player.username]) {{
                    const mesh = new THREE.Mesh(playerGeometry, new THREE.MeshPhongMaterial({{ color: player.color }}));
                    mesh.castShadow = true;
                    scene.add(mesh);
                    playerMeshes[player.username] = mesh;
                }}
                const mesh = playerMeshes[player.username];
                mesh.position.set(player.x, 1, player.z);
                players[player.username] = {{ mesh: mesh, score: player.score, treasures: player.treasures }};
                if (player.username === playerName) {{
                    xPos = player.x;
                    zPos = player.z;
                    score = player.score;
                    treasureCount = player.treasures;
                    playerMesh.position.set(xPos, 1, zPos);
                }}
            }});
            document.getElementById('players').textContent = `Players: ${{Object.keys(playerMeshes).length}}`;
            document.getElementById('score').textContent = `Score: ${{score}}`;
            document.getElementById('treasures').textContent = `Treasures: ${{treasureCount}}`;
        }}

        ws.onmessage = function(event) {{
            const data = event.data;
            if (data.startsWith('MAP_UPDATE:')) {{
                const playerData = JSON.parse(data.split('MAP_UPDATE:')[1]);
                updatePlayers(playerData);
            }} else if (data.startsWith('CHAT_UPDATE:')) {{
                const chatData = JSON.parse(data.split('CHAT_UPDATE:')[1]);
                const chatBox = document.getElementById('chatBox');
                chatBox.innerHTML = chatData.map(line => `<p>${{line}}</p>`).join('');
                chatBox.scrollTop = chatBox.scrollHeight;
            }} else if (data.startsWith('GAME_STATE:')) {{
                const gameState = JSON.parse(data.split('GAME_STATE:')[1]);
                updatePlayers(gameState.players);
                updateTreasures(gameState.treasures);
            }} else if (!data.startsWith('PRAIRIE_UPDATE:')) {{
                const [sender, message] = data.split('|');
                const chatBox = document.getElementById('chatBox');
                chatBox.innerHTML += `<p>${{sender}}: ${{message}}</p>`;
                chatBox.scrollTop = chatBox.scrollHeight;
            }}
        }};

        let lastTime = performance.now();
        function animate() {{
            requestAnimationFrame(animate);
            const currentTime = performance.now();
            const delta = (currentTime - lastTime) / 1000;
            lastTime = currentTime;

            updatePlayer(delta);
            treasures.forEach(t => t.rotation.y += delta);
            renderer.render(scene, camera);
        }}

        animate();
    </script>
</body>
</html>
"""

# Main Game Loop
def main():
    st.sidebar.title(f"🎮 {GAME_NAME}")
    st.sidebar.subheader(f"🌄 {st.session_state.username}’s Adventure - Score: {st.session_state.score} 🏆")
    st.sidebar.write(f"📜 {next(c['desc'] for c in CHARACTERS if c['name'] == st.session_state.username)}")
    st.sidebar.write(f"📍 Location: {st.session_state.location}")
    st.sidebar.write(f"🏅 Score: {st.session_state.score}")
    st.sidebar.write(f"🎵 Treasures: {st.session_state.treasures}")
    st.sidebar.write(f"👥 Players: {', '.join([p['username'] for p in st.session_state.players.values()]) or 'None'}")
    
    st.session_state.update_interval = st.sidebar.slider("Refresh Interval (seconds)", 1, 10, 1, step=1)
    if st.sidebar.button("Reset World"):
        reset_game_state()
        st.session_state.game_state_timestamp = time.time()
        st.rerun()

    left_col, right_col = st.columns([2, 1])

    with left_col:
        components.html(rocky_map_html, width=800, height=600)
        chat_content = asyncio.run(load_chat())
        st.text_area("📜 Quest Log", "\n".join(chat_content[-10:]), height=200, disabled=True)
        message = st.text_input(f"🗨️ {st.session_state.username} says:", placeholder="Speak or type to chat! 🌲")
        if st.button("🌟 Send & Chat 🎤"):
            if message:
                voice = next(c["voice"] for c in CHARACTERS if c["name"] == st.session_state.username)
                md_file, audio_file = asyncio.run(save_chat_entry(st.session_state.username, message, voice))
                if audio_file:
                    play_and_download_audio(audio_file)
                st.session_state.last_activity = time.time()
                st.success(f"🌄 +10 points! New Score: {st.session_state.score}")
        mycomponent = components.declare_component("speech_component", path="./speech_component")
        val = mycomponent(my_input_value="", key=f"speech_{st.session_state.get('speech_processed', False)}")
        if val and val != st.session_state.last_transcript:
            val_stripped = val.strip().replace('\n', ' ')
            if val_stripped:
                voice = next(c["voice"] for c in CHARACTERS if c["name"] == st.session_state.username)
                st.session_state['speech_processed'] = True
                md_file, audio_file = asyncio.run(save_chat_entry(st.session_state.username, val_stripped, voice))
                if audio_file:
                    play_and_download_audio(audio_file)
                st.session_state.last_activity = time.time()
                st.rerun()

    with right_col:
        st.subheader("🌾 Prairie Map")
        prairie_map = folium.Map(location=[44.0, -103.0], zoom_start=8, tiles="CartoDB Positron")
        for loc, (lat, lon) in PRAIRIE_LOCATIONS.items():
            folium.Marker([lat, lon], popup=loc).add_to(prairie_map)
        game_state = load_game_state(st.session_state.game_state_timestamp)
        for player in game_state["players"].values():
            folium.CircleMarker(
                location=[44.0 + player["x"] * 0.01, -103.0 + player["z"] * 0.01],  # Approximate mapping to lat/lon
                radius=5,
                color=f"#{player['color']:06x}",
                fill=True,
                fill_opacity=0.7,
                popup=f"{player['username']} (Score: {player['score']}, Treasures: {player['treasures']})"
            ).add_to(prairie_map)
        for client_id, player in st.session_state.prairie_players.items():
            folium.CircleMarker(
                location=player['location'],
                radius=5,
                color=f"#{player['color']:06x}",
                fill=True,
                fill_opacity=0.7,
                popup=f"{player['username']} ({player['animal']})"
            ).add_to(prairie_map)
        folium_static(prairie_map, width=600, height=400)
        
        animal = st.selectbox("Choose Animal", ["prairie_dog", "deer", "sheep", "groundhog"])
        location = st.selectbox("Move to", list(PRAIRIE_LOCATIONS.keys()))
        if st.button("Move"):
            asyncio.run(broadcast_message(f"{st.session_state.username}|PRAIRIE:ANIMAL:{animal}", "quest"))
            asyncio.run(broadcast_message(f"{st.session_state.username}|PRAIRIE:MOVE:{location}", "quest"))
            st.session_state.last_activity = time.time()
            st.rerun()

    elapsed = time.time() - st.session_state.last_update
    remaining = max(0, st.session_state.update_interval - elapsed)
    st.sidebar.markdown(f"⏳ Next Update in: {int(remaining)}s")
    
    # Timeout and Auto-Refresh Logic
    if time.time() - st.session_state.last_activity > st.session_state.timeout:
        st.sidebar.warning("Timed out! Refreshing in 5 seconds...")
        time.sleep(5)
        st.session_state.game_state_timestamp = time.time()
        st.rerun()
    elif time.time() - st.session_state.last_update > st.session_state.auto_refresh:
        st.session_state.game_state_timestamp = time.time()
        st.rerun()
    elif remaining <= 0:
        st.session_state.last_update = time.time()
        st.session_state.game_state_timestamp = time.time()
        st.rerun()

    if not st.session_state.get('server_running', False):
        st.session_state.server_task = threading.Thread(target=start_websocket_server, daemon=True)
        st.session_state.server_task.start()

if __name__ == "__main__":
    main()