import streamlit as st
import streamlit.components.v1 as components
import asyncio
import websockets
import uuid
import os
import random
import time
import hashlib
from datetime import datetime
import pytz
import nest_asyncio
import edge_tts
from audio_recorder_streamlit import audio_recorder

# Patch asyncio for nesting
nest_asyncio.apply()

# Page Config
st.set_page_config(page_title="Galaxian Snake 3D Multiplayer", layout="wide")

st.title("Galaxian Snake 3D Multiplayer")
st.write("Navigate a 3D city, grow your snake, shoot enemies, and chat!")

# Sliders for container size
max_width = min(1200, st.session_state.get('window_width', 1200))
max_height = min(1600, st.session_state.get('window_height', 1600))

col1, col2 = st.columns(2)
with col1:
    container_width = st.slider("Container Width (px)", 300, max_width, 768, step=50)
with col2:
    container_height = st.slider("Container Height (px)", 400, max_height, 1024, step=50)

# Session State Initialization
def init_session_state():
    defaults = {
        'server_running': False, 'active_connections': {},
        'chat_history': [], 'last_chat_update': 0, 'username': None,
        'tts_voice': "en-US-AriaNeural", 'audio_cache': {}
    }
    for k, v in defaults.items():
        if k not in st.session_state:
            st.session_state[k] = v

init_session_state()

# Usernames and Voices
FUN_USERNAMES = {
    "CosmicJester 🌌": "en-US-AriaNeural",
    "PixelPanda 🐼": "en-US-JennyNeural",
    "QuantumQuack 🦆": "en-GB-SoniaNeural",
    "StellarSquirrel 🐿️": "en-AU-NatashaNeural",
    "GizmoGuru ⚙️": "en-CA-ClaraNeural",
    "NebulaNinja 🌠": "en-US-GuyNeural",
}

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]

# Chat File
CHAT_FILE = "chat_logs/global_chat.md"
os.makedirs("chat_logs", exist_ok=True)
if not os.path.exists(CHAT_FILE):
    with open(CHAT_FILE, 'a') as f:
        f.write("# Multiplayer Snake Chat\n\nWelcome to the cosmic city! 🎤\n")

# 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]
    filename = f"audio_logs/{datetime.now().strftime('%Y%m%d_%H%M%S')}-by-{username}-{hashlib.md5(text.encode()).hexdigest()[:8]}.mp3"
    os.makedirs("audio_logs", exist_ok=True)
    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):
        with open(file_path, "rb") as f:
            audio_bytes = f.read()
        st.audio(audio_bytes, format="audio/mp3")
        b64 = base64.b64encode(audio_bytes).decode()
        st.markdown(f'<a href="data:audio/mpeg;base64,{b64}" download="{os.path.basename(file_path)}">🎵 Download {os.path.basename(file_path)}</a>', unsafe_allow_html=True)

# WebSocket Handling
async def websocket_handler(websocket, path):
    client_id = str(uuid.uuid4())
    room_id = "snake_chat"
    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
    await broadcast_message(f"System|{username} has joined the game!", room_id)
    try:
        async for message in websocket:
            if '|' in message:
                sender, content = message.split('|', 1)
                await save_chat_entry(sender, content)
            else:
                await websocket.send("ERROR|Message format: username|content")
    except websockets.ConnectionClosed:
        await broadcast_message(f"System|{username} has left the game!", room_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 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]

async def start_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
        st.session_state['server'] = server
        await asyncio.Future()

# Chat Functions
async def save_chat_entry(username, message):
    central = pytz.timezone('US/Central')
    timestamp = datetime.now(central).strftime("%Y-%m-%d %H:%M:%S")
    entry = f"[{timestamp}] {username}: {message}"
    with open(CHAT_FILE, 'a') as f:
        f.write(f"{entry}\n")
    audio_file = await async_edge_tts_generate(message, FUN_USERNAMES.get(username, "en-US-AriaNeural"), username)
    if audio_file:
        play_and_download_audio(audio_file)
    await broadcast_message(f"{username}|{message}", "snake_chat")
    st.session_state.chat_history.append(entry)
    st.session_state.last_chat_update = time.time()

async def load_chat():
    with open(CHAT_FILE, 'r') as f:
        content = f.read().strip()
    return content.split('\n')

# Game HTML
html_code = f"""
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Galaxian Snake 3D Multiplayer</title>
    <style>
        body {{ margin: 0; overflow: hidden; font-family: Arial, sans-serif; background: #000; }}
        #gameContainer {{ width: {container_width}px; height: {container_height}px; position: relative; }}
        canvas {{ width: 100%; height: 100%; display: block; }}
        .ui-container {{
            position: absolute; top: 10px; left: 10px; color: white;
            background-color: rgba(0, 0, 0, 0.5); padding: 10px; border-radius: 5px;
            user-select: none;
        }}
        .controls {{
            position: absolute; bottom: 10px; left: 10px; color: white;
            background-color: rgba(0, 0, 0, 0.5); padding: 10px; border-radius: 5px;
        }}
        #chatBox {{ 
            position: absolute; bottom: 60px; left: 10px; width: 300px; height: 200px; 
            background: rgba(0, 0, 0, 0.7); color: white; padding: 10px; 
            border-radius: 5px; overflow-y: auto; 
        }}
        #debug {{ position: absolute; top: 10px; right: 10px; color: white; }}
    </style>
</head>
<body>
    <div id="gameContainer">
        <div class="ui-container">
            <h2>Galaxian Snake 3D</h2>
            <div id="players">Players: 1</div>
            <div id="score">Score: 0</div>
            <div id="length">Length: 1</div>
            <div id="survival">Survival: 0s</div>
            <div id="shield">Shield: Off</div>
        </div>
        <div id="chatBox"></div>
        <div class="controls">
            <p>Controls: W/A/S/D or Arrows to steer, Space to shoot, Shift for shield</p>
            <p>Eat food, shoot buildings, pass gates, dodge enemies!</p>
        </div>
        <div id="debug"></div>
    </div>

    <script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script>
    <script>
        try {{
            console.log('Starting game initialization...');
            let score = 0, players = {{}}, foodItems = [], buildings = [], bullets = [], debris = [], gates = [], enemies = [], enemyBullets = [];
            let snakeSegments = [], moveLeft = false, moveRight = false, moveUp = false, moveDown = false, shoot = false, shield = false;
            const playerName = "{st.session_state.username}";
            let ws = new WebSocket('ws://localhost:8765');
            const debug = document.getElementById('debug');
            let survivalTime = 0, lastCollisionTime = performance.now(), shieldTimer = 0;

            // Scene setup
            const scene = new THREE.Scene();
            const camera = new THREE.PerspectiveCamera(75, {container_width} / {container_height}, 0.1, 1000);
            camera.position.set(0, 15, 20);
            console.log('Camera initialized at:', camera.position);

            const renderer = new THREE.WebGLRenderer({{ antialias: true }});
            renderer.setSize({container_width}, {container_height});
            renderer.shadowMap.enabled = true;
            document.getElementById('gameContainer').appendChild(renderer.domElement);
            console.log('Renderer initialized');

            // Lighting
            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);
            console.log('Lights added');

            // Ground
            const textureLoader = new THREE.TextureLoader();
            const groundGeometry = new THREE.PlaneGeometry(200, 1000);
            const groundMaterial = new THREE.MeshStandardMaterial({{ 
                color: 0x1a5e1a,
                bumpMap: textureLoader.load('https://threejs.org/examples/textures/terrain/grasslight-big-nm.jpg'),
                bumpScale: 0.1
            }});
            const ground = new THREE.Mesh(groundGeometry, groundMaterial);
            ground.rotation.x = -Math.PI / 2;
            ground.receiveShadow = true;
            ground.position.z = -500;
            scene.add(ground);
            console.log('Ground added');

            // Player (Snake Head and Segments)
            const snakeGeometry = new THREE.BoxGeometry(1, 1, 1);
            const snakeMaterial = new THREE.MeshPhongMaterial({{ color: 0x00ff00 }});
            const head = new THREE.Mesh(snakeGeometry, snakeMaterial);
            head.position.set(0, 0.5, 0);
            head.castShadow = true;
            snakeSegments.push({{ mesh: head, timer: -1, permanent: true }}); // Head is permanent
            scene.add(head);
            players[playerName] = {{ snake: snakeSegments, score: 0, length: 1 }};
            console.log('Snake head initialized at:', head.position);

            // Building rules
            const buildingColors = [0x888888, 0x666666, 0x999999];
            const buildingRules = [
                {{name: "Simple", axiom: "F", rules: {{"F": "F[+F][-F]"}}, iterations: 1, baseHeight: 10, baseWidth: 5, baseDepth: 5, angle: Math.PI/4, probability: 1}}
            ];

            function interpretLSystem(rule, position, rotation) {{
                let currentString = rule.axiom;
                for (let i = 0; i < rule.iterations; i++) {{
                    let newString = "";
                    for (let j = 0; j < currentString.length; j++) {{
                        newString += rule.rules[currentString[j]] || currentString[j];
                    }}
                    currentString = newString;
                }}
                
                let buildingGroup = new THREE.Group();
                buildingGroup.position.copy(position);
                const stack = [];
                let currentPosition = new THREE.Vector3(0, 0, 0);
                let currentRotation = rotation || new THREE.Euler();
                let scale = new THREE.Vector3(1, 1, 1);
                const color = buildingColors[Math.floor(Math.random() * buildingColors.length)];
                const material = new THREE.MeshStandardMaterial({{color: color}});
                
                for (let i = 0; i < currentString.length; i++) {{
                    const char = currentString[i];
                    switch (char) {{
                        case 'F':
                            const width = rule.baseWidth * scale.x;
                            const height = rule.baseHeight * scale.y;
                            const depth = rule.baseDepth * scale.z;
                            const geometry = new THREE.BoxGeometry(width, height, depth);
                            const buildingPart = new THREE.Mesh(geometry, material);
                            buildingPart.position.copy(currentPosition);
                            buildingPart.position.y += height / 2;
                            buildingPart.rotation.copy(currentRotation);
                            buildingPart.castShadow = true;
                            buildingPart.receiveShadow = true;
                            buildingGroup.add(buildingPart);
                            const direction = new THREE.Vector3(0, height, 0);
                            direction.applyEuler(currentRotation);
                            currentPosition.add(direction);
                            break;
                        case '+': currentRotation.y += rule.angle; break;
                        case '-': currentRotation.y -= rule.angle; break;
                        case '[': stack.push({{position: currentPosition.clone(), rotation: currentRotation.clone(), scale: scale.clone()}}); break;
                        case ']': if (stack.length > 0) {{ const state = stack.pop(); currentPosition = state.position; currentRotation = state.rotation; scale = state.scale; }} break;
                    }}
                }}
                return buildingGroup;
            }}

            function spawnBuildings() {{
                const cityWidth = 40, cityDepth = -1000;
                for (let i = buildings.length; i < 20; i++) {{
                    const position = new THREE.Vector3(
                        Math.random() * cityWidth - cityWidth / 2,
                        0,
                        cityDepth + Math.random() * 1000
                    );
                    const building = interpretLSystem(buildingRules[0], position, new THREE.Euler());
                    buildings.push(building);
                    scene.add(building);
                    console.log('Building at:', position);
                }}
            }}

            function spawnFood() {{
                if (foodItems.length < 10) {{
                    const food = new THREE.Mesh(
                        new THREE.BoxGeometry(1, 1, 1),
                        new THREE.MeshStandardMaterial({{color: 0xffff00}})
                    );
                    food.position.set(
                        Math.random() * 40 - 20,
                        Math.random() * 10,
                        -1000
                    );
                    foodItems.push(food);
                    scene.add(food);
                    console.log('Food at:', food.position);
                }}
            }}

            function spawnGates() {{
                if (gates.length < 5) {{
                    const gate = new THREE.Group();
                    const postGeometry = new THREE.CylinderGeometry(0.5, 0.5, 10, 8);
                    const postMaterial = new THREE.MeshPhongMaterial({{ color: 0xaaaaaa }});
                    const leftPost = new THREE.Mesh(postGeometry, postMaterial);
                    leftPost.position.set(-5, 5, 0);
                    const rightPost = new THREE.Mesh(postGeometry, postMaterial);
                    rightPost.position.set(5, 5, 0);
                    const fieldGeometry = new THREE.PlaneGeometry(10, 10);
                    const fieldMaterial = new THREE.MeshBasicMaterial({{ color: 0x00ffff, transparent: true, opacity: 0.5, side: THREE.DoubleSide }});
                    const field = new THREE.Mesh(fieldGeometry, fieldMaterial);
                    field.position.set(0, 5, 0);
                    gate.add(leftPost);
                    gate.add(rightPost);
                    gate.add(field);
                    gate.position.set(Math.random() * 40 - 20, 0, -1000 + Math.random() * 1000);
                    gates.push(gate);
                    scene.add(gate);
                    console.log('Gate at:', gate.position);
                }}
            }}

            function spawnEnemyCluster() {{
                const clusterCenter = new THREE.Vector3(
                    Math.random() * 40 - 20,
                    Math.random() * 10 + 5,
                    -1000
                );
                const enemyGeometry = new THREE.BoxGeometry(0.8, 0.8, 0.8);
                const enemyMaterial = new THREE.MeshPhongMaterial({{ color: 0xff0000 }});
                const pattern = [[-2, 1], [0, 1], [2, 1], [-1, 2], [1, 2]];
                pattern.forEach(pos => {{
                    const enemy = new THREE.Mesh(enemyGeometry, enemyMaterial);
                    enemy.position.set(
                        clusterCenter.x + pos[0],
                        clusterCenter.y + pos[1],
                        clusterCenter.z
                    );
                    enemy.shootTimer = Math.random();
                    enemies.push(enemy);
                    scene.add(enemy);
                }});
                console.log('Enemy cluster at:', clusterCenter);
            }}

            function addTailSegment() {{
                if (snakeSegments.length < 5) {{
                    const segment = new THREE.Mesh(snakeGeometry, snakeMaterial.clone());
                    const lastSegment = snakeSegments[snakeSegments.length - 1].mesh;
                    segment.position.copy(lastSegment.position);
                    segment.castShadow = true;
                    snakeSegments.push({{ mesh: segment, timer: 5, permanent: false }});
                    scene.add(segment);
                    players[playerName].length = snakeSegments.length;
                    console.log('Tail segment added, length:', snakeSegments.length);
                }}
            }}

            function shootBullet(segment) {{
                const bulletGeometry = new THREE.SphereGeometry(0.3, 8, 8);
                const bulletMaterial = new THREE.MeshPhongMaterial({{ color: 0xff0000 }});
                const bullet = new THREE.Mesh(bulletGeometry, bulletMaterial);
                bullet.position.copy(segment.position);
                bullet.velocity = new THREE.Vector3(0, 0, -20);
                bullet.timer = 5;
                bullets.push(bullet);
                scene.add(bullet);
            }}

            function spawnExplosionParticles(position) {{
                const particleGeometry = new THREE.SphereGeometry(0.2, 8, 8);
                const particleMaterial = new THREE.MeshBasicMaterial({{ color: 0xffff00 }});
                for (let i = 0; i < 20; i++) {{
                    const particle = new THREE.Mesh(particleGeometry, particleMaterial);
                    particle.position.copy(position);
                    particle.velocity = new THREE.Vector3(
                        Math.random() - 0.5,
                        Math.random() - 0.5,
                        Math.random() - 0.5
                    ).multiplyScalar(5);
                    particle.timer = 1;
                    debris.push(particle);
                    scene.add(particle);
                }}
            }}

            function destroyBuilding(building, index) {{
                building.children.forEach(child => {{
                    const debrisPiece = new THREE.Mesh(child.geometry, child.material);
                    debrisPiece.position.copy(child.getWorldPosition(new THREE.Vector3()));
                    debrisPiece.velocity = new THREE.Vector3(
                        Math.random() - 0.5,
                        Math.random() * 2,
                        Math.random() - 0.5
                    ).multiplyScalar(5);
                    debrisPiece.timer = 3;
                    debris.push(debrisPiece);
                    scene.add(debrisPiece);
                }});
                scene.remove(building);
                buildings.splice(index, 1);
                score += 100;
                spawnBuildings();
            }}

            function shootEnemyBullet(enemy) {{
                const bulletGeometry = new THREE.SphereGeometry(0.3, 8, 8);
                const bulletMaterial = new THREE.MeshPhongMaterial({{ color: 0x00ff00 }});
                const bullet = new THREE.Mesh(bulletGeometry, bulletMaterial);
                bullet.position.copy(enemy.position);
                const direction = snakeSegments[0].mesh.position.clone().sub(enemy.position).normalize();
                bullet.velocity = direction.multiplyScalar(15);
                bullet.timer = 5;
                enemyBullets.push(bullet);
                scene.add(bullet);
            }}

            // 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': shoot = true; break;
                    case 'ShiftLeft': case 'ShiftRight': shield = true; shieldTimer = 4; break;
                }}
            }});
            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': shoot = false; break;
                }}
            }});

            function updatePlayer(delta) {{
                const speed = 10;
                const head = snakeSegments[0].mesh;
                if (moveLeft && head.position.x > -20) head.position.x -= speed * delta;
                if (moveRight && head.position.x < 20) head.position.x += speed * delta;
                if (moveUp && head.position.y < 15) head.position.y += speed * delta;
                if (moveDown && head.position.y > 0.5) head.position.y -= speed * delta;

                if (shoot && !bullets.some(b => b.timer > 4.8)) {{
                    snakeSegments.forEach(segment => shootBullet(segment.mesh));
                }}
                if (shieldTimer > 0) {{
                    shieldTimer -= delta;
                    snakeSegments.forEach(s => s.mesh.material.color.setHex(0x0000ff));
                }} else {{
                    snakeSegments.forEach(s => s.mesh.material.color.setHex(0x00ff00));
                }}

                // Update tail positions
                for (let i = snakeSegments.length - 1; i > 0; i--) {{
                    const current = snakeSegments[i].mesh;
                    const previous = snakeSegments[i - 1].mesh;
                    current.position.lerp(previous.position, 0.5);
                    if (!snakeSegments[i].permanent) {{
                        snakeSegments[i].timer -= delta;
                        if (snakeSegments[i].timer <= 0) {{
                            scene.remove(current);
                            snakeSegments.splice(i, 1);
                            players[playerName].length = snakeSegments.length;
                        }}
                    }}
                }}

                const forwardSpeed = 20;
                buildings.forEach((building, i) => {{
                    building.position.z += forwardSpeed * delta;
                    if (building.position.z > 20) {{
                        building.position.z = -1000;
                        building.position.x = Math.random() * 40 - 20;
                    }}
                    if (head.position.distanceTo(building.position) < 5) {{
                        head.position.set(0, 0.5, 0);
                        lastCollisionTime = performance.now();
                    }}
                }});
                foodItems.forEach((food, i) => {{
                    food.position.z += forwardSpeed * delta;
                    if (food.position.z > 20) {{
                        scene.remove(food);
                        foodItems.splice(i, 1);
                        spawnFood();
                    }}
                    if (head.position.distanceTo(food.position) < 1) {{
                        scene.remove(food);
                        foodItems.splice(i, 1);
                        score += 10;
                        players[playerName].length = snakeSegments.length;
                        spawnFood();
                    }}
                }});
                gates.forEach((gate, i) => {{
                    gate.position.z += forwardSpeed * delta;
                    if (gate.position.z > 20) {{
                        scene.remove(gate);
                        gates.splice(i, 1);
                        spawnGates();
                    }}
                    const fieldPos = gate.children[2].getWorldPosition(new THREE.Vector3());
                    if (head.position.distanceTo(fieldPos) < 5) {{
                        score += 100;
                        addTailSegment();
                        for (let j = 0; j < snakeSegments.length; j++) {{
                            if (!snakeSegments[j].permanent && snakeSegments[j].timer > 0) {{
                                snakeSegments[j].permanent = true;
                                snakeSegments[j].timer = -1;
                                score += 20;
                                spawnExplosionParticles(fieldPos);
                                break;
                            }}
                        }}
                        scene.remove(gate);
                        gates.splice(i, 1);
                        spawnGates();
                    }}
                }});

                camera.position.set(head.position.x, head.position.y + 15, head.position.z + 20);
                camera.lookAt(head.position);
            }}

            function updateBullets(delta) {{
                for (let i = bullets.length - 1; i >= 0; i--) {{
                    const bullet = bullets[i];
                    bullet.position.add(bullet.velocity.clone().multiplyScalar(delta));
                    bullet.timer -= delta;

                    if (bullet.timer <= 0) {{
                        scene.remove(bullet);
                        bullets.splice(i, 1);
                        continue;
                    }}

                    for (let j = buildings.length - 1; j >= 0; j--) {{
                        if (bullet.position.distanceTo(buildings[j].position) < 5) {{
                            destroyBuilding(buildings[j], j);
                            scene.remove(bullet);
                            bullets.splice(i, 1);
                            break;
                        }}
                    }}
                }}
            }}

            function updateDebris(delta) {{
                for (let i = debris.length - 1; i >= 0; i--) {{
                    const piece = debris[i];
                    piece.position.add(piece.velocity.clone().multiplyScalar(delta));
                    piece.velocity.y -= 9.8 * delta;
                    piece.timer -= delta;

                    if (piece.position.y <= 0) {{
                        piece.velocity.y *= -0.7;
                        piece.position.y = 0;
                    }}
                    if (piece.timer <= 0) {{
                        scene.remove(piece);
                        debris.splice(i, 1);
                    }}
                }}
            }}

            function updateEnemies(delta) {{
                enemies.forEach((enemy, i) => {{
                    enemy.position.z += 10 * delta;
                    enemy.shootTimer -= delta;
                    if (enemy.shootTimer <= 0) {{
                        shootEnemyBullet(enemy);
                        enemy.shootTimer = 1;
                    }}
                    if (enemy.position.z > 20) {{
                        scene.remove(enemy);
                        enemies.splice(i, 1);
                    }}
                }});
                if (enemies.length < 10) spawnEnemyCluster();
            }}

            function updateEnemyBullets(delta) {{
                for (let i = enemyBullets.length - 1; i >= 0; i--) {{
                    const bullet = enemyBullets[i];
                    const direction = snakeSegments[0].mesh.position.clone().sub(bullet.position).normalize();
                    bullet.velocity.lerp(direction.multiplyScalar(15), delta * 2);
                    bullet.position.add(bullet.velocity.clone().multiplyScalar(delta));
                    bullet.timer -= delta;

                    if (bullet.timer <= 0 || bullet.position.z > 20) {{
                        scene.remove(bullet);
                        enemyBullets.splice(i, 1);
                        continue;
                    }}

                    if (bullet.position.distanceTo(snakeSegments[0].mesh.position) < 1 && shieldTimer <= 0) {{
                        snakeSegments[0].mesh.position.set(0, 0.5, 0);
                        lastCollisionTime = performance.now();
                        scene.remove(bullet);
                        enemyBullets.splice(i, 1);
                    }} else if (bullet.position.distanceTo(snakeSegments[0].mesh.position) < 1 && shieldTimer > 0) {{
                        scene.remove(bullet);
                        enemyBullets.splice(i, 1);
                    }}
                }}
            }}

            function updateUI() {{
                survivalTime = (performance.now() - lastCollisionTime) / 1000;
                score += Math.floor(survivalTime / 10);
                document.getElementById('players').textContent = `Players: ${{Object.keys(players).length}}`;
                document.getElementById('score').textContent = `Score: ${{score}}`;
                document.getElementById('length').textContent = `Length: ${{snakeSegments.length}}`;
                document.getElementById('survival').textContent = `Survival: ${{survivalTime.toFixed(1)}}s`;
                document.getElementById('shield').textContent = `Shield: ${{shieldTimer > 0 ? 'On' : 'Off'}}`;
                debug.innerHTML = `Scene: ${{scene.children.length}}<br>Buildings: ${{buildings.length}}<br>Food: ${{foodItems.length}}<br>Bullets: ${{bullets.length}}<br>Gates: ${{gates.length}}<br>Enemies: ${{enemies.length}}`;
            }}

            // WebSocket chat
            ws.onmessage = function(event) {{
                const [sender, message] = event.data.split('|');
                const chatBox = document.getElementById('chatBox');
                chatBox.innerHTML += `<p>${{sender}}: ${{message}}</p>`;
                chatBox.scrollTop = chatBox.scrollHeight;
            }};

            ws.onopen = function() {{ console.log('Connected to WebSocket server'); }};
            ws.onerror = function(error) {{ console.error('WebSocket error:', error); }};
            ws.onclose = function() {{ console.log('Disconnected from WebSocket server'); }};

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

                updatePlayer(delta);
                updateBullets(delta);
                updateDebris(delta);
                updateEnemies(delta);
                updateEnemyBullets(delta);
                spawnFood();
                spawnGates();
                updateUI();

                renderer.render(scene, camera);
            }}

            // Initialize
            spawnBuildings();
            spawnFood();
            spawnGates();
            spawnEnemyCluster();
            animate();
            console.log('Game initialized');
        }} catch (e) {{
            console.error('Error in game initialization:', e);
            document.getElementById('debug').innerHTML = 'Error: ' + e.message;
        }}
    </script>
</body>
</html>
"""

# Render the HTML component
components.html(html_code, width=container_width, height=container_height)

# Chat Interface
st.sidebar.title(f"Chat as {st.session_state.username}")
chat_content = asyncio.run(load_chat())
chat_container = st.sidebar.container()
with chat_container:
    st.code("\n".join(chat_content), language="python")

message = st.sidebar.text_input("Message", key="chat_input")
if message and st.sidebar.button("Send 🚀"):
    asyncio.run(save_chat_entry(st.session_state.username, message))

audio_bytes = audio_recorder()
if audio_bytes:
    with open("temp_audio.wav", "wb") as f:
        f.write(audio_bytes)
    message = "Voice message received"
    asyncio.run(save_chat_entry(st.session_state.username, message))

# Start WebSocket Server in Main Event Loop
async def main_async():
    if not st.session_state.get('server_running', False):
        await start_websocket_server()

loop = asyncio.get_event_loop()
if not st.session_state.get('server_running', False):
    loop.run_until_complete(main_async())

st.sidebar.write("""
### How to Play
- **W/A/S/D or Arrow Keys**: Steer the snake (up/down/left/right)
- **Space**: Shoot bullets from all segments to destroy buildings
- **Shift**: Activate 4s shield to block enemy bullets
- Eat yellow cubes for 10 points
- Destroy buildings for 100 points each
- Pass through gates for 100 points + tail segment (up to 5)
- Cross gate again within 5s for permanent segment + 20 points + explosion
- Survive longer for bonus points (1 per 10s)
- Dodge heat-seeking enemy bullets
- Chat with other players in real-time
""")