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 with continuous motion, eat food, and chat in real-time!")

# 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: 3</div>
        </div>
        <div id="chatBox"></div>
        <div class="controls">
            <p>Controls: W/A/S/D or Arrow Keys to steer</p>
            <p>Eat yellow cubes to grow and score!</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 = [];
            let snake, moveLeft = false, moveRight = false, moveUp = false, moveDown = false;
            const playerName = "{st.session_state.username}";
            let ws = new WebSocket('ws://localhost:8765');
            const debug = document.getElementById('debug');

            // 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); // Long scrolling ground
            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; // Start far back
            scene.add(ground);
            console.log('Ground added');

            // Player (Snake Head)
            const snakeGeometry = new THREE.BoxGeometry(1, 1, 1);
            const snakeMaterial = new THREE.MeshPhongMaterial({{ color: 0x00ff00 }});
            snake = new THREE.Mesh(snakeGeometry, snakeMaterial);
            snake.position.set(0, 0.5, 0);
            snake.castShadow = true;
            scene.add(snake);
            players[playerName] = {{ snake: snake, score: 0, length: 1 }};
            console.log('Snake initialized at:', snake.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);
                }}
            }}

            // 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;
                }}
            }});
            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;
                }}
            }});

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

                // Continuous forward motion
                const forwardSpeed = 20;
                buildings.forEach(building => {{
                    building.position.z += forwardSpeed * delta;
                    if (building.position.z > 20) {{
                        building.position.z = -1000;
                        building.position.x = Math.random() * 40 - 20;
                    }}
                    if (snake.position.distanceTo(building.position) < 5) {{
                        snake.position.set(0, 0.5, 0); // Reset on collision
                    }}
                }});
                foodItems.forEach((food, i) => {{
                    food.position.z += forwardSpeed * delta;
                    if (food.position.z > 20) {{
                        scene.remove(food);
                        foodItems.splice(i, 1);
                        spawnFood();
                    }}
                    if (snake.position.distanceTo(food.position) < 1) {{
                        scene.remove(food);
                        foodItems.splice(i, 1);
                        score += 10;
                        players[playerName].length += 1;
                        spawnFood();
                    }}
                }});

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

            function updateUI() {{
                document.getElementById('players').textContent = `Players: ${{Object.keys(players).length}}`;
                document.getElementById('score').textContent = `Score: ${{score}}`;
                document.getElementById('length').textContent = `Length: ${{players[playerName].length}}`;
                debug.innerHTML = `Scene: ${{scene.children.length}}<br>Buildings: ${{buildings.length}}<br>Food: ${{foodItems.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);
                spawnFood();
                updateUI();

                renderer.render(scene, camera);
            }}

            // Initialize
            spawnBuildings();
            spawnFood();
            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)
- Eat yellow cubes to grow and score
- World moves continuously toward you
- Chat with other players in real-time
""")