import streamlit as st
import asyncio
import websockets
import uuid
from datetime import datetime
import os
import random
import hashlib
import glob
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

# Patch asyncio for nesting
nest_asyncio.apply()

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

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

# Directories
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")

# 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"{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 (Moved Up for save_chat_entry Dependency)
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.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
    st.session_state.score += 10
    st.session_state.treasures += 1
    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 (Now After All Dependencies)
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': 20, 'x_pos': 0, 'z_pos': 0, 'move_left': False,
        'move_right': False, 'move_up': False, 'move_down': False
    }
    for k, v in defaults.items():
        if k not in st.session_state:
            st.session_state[k] = v
    # Ensure username is initialized immediately
    if st.session_state.username is None:
        saved_username = load_username()
        if saved_username and saved_username in CHARACTERS:
            st.session_state.username = saved_username
        else:
            st.session_state.username = random.choice(list(CHARACTERS.keys()))
            asyncio.run(save_chat_entry(st.session_state.username, "🗺️ Begins the Rocky Mountain Quest!", CHARACTERS[st.session_state.username]["voice"]))
        save_username(st.session_state.username)

# Call init_session_state immediately to ensure username is set
init_session_state()

# 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 = CHARACTERS[username]["voice"]
    md_file, audio_file = await save_chat_entry(username, result, voice, True)
    return md_file, audio_file

# WebSocket for Multiplayer with Map Updates
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
    st.session_state.players[client_id] = {
        "username": username,
        "x": random.uniform(-20, 20),
        "z": random.uniform(-50, 50),
        "color": CHARACTERS[username]["color"]
    }
    await save_chat_entry(username, f"🗺️ Joins the quest at {START_LOCATION}!", CHARACTERS[username]["voice"])
    
    try:
        async for message in websocket:
            if '|' in message:
                username, content = message.split('|', 1)
                voice = CHARACTERS.get(username, {"voice": "en-US-AriaNeural"})["voice"]
                if content.startswith("MOVE:"):
                    _, x, z = content.split(":")
                    st.session_state.players[client_id]["x"] = float(x)
                    st.session_state.players[client_id]["z"] = float(z)
                else:
                    await save_chat_entry(username, content, voice)
                    await perform_arxiv_search(content, username)
    except websockets.ConnectionClosed:
        await save_chat_entry(username, "🏃 Leaves the quest!", CHARACTERS[username]["voice"])
        if client_id in st.session_state.players:
            del st.session_state.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"):
            player_list = ", ".join([p["username"] for p in st.session_state.players.values()]) or "No adventurers yet!"
            message = f"📢 Quest Update: Active Adventurers - {player_list}"
            player_data = json.dumps(list(st.session_state.players.values()))
            await broadcast_message(f"System|{message}", "quest")
            await broadcast_message(f"MAP_UPDATE:{player_data}", "quest")
            await save_chat_entry("System", 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())

# Game HTML with Map
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>Rocky Mountain Quest Map</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; }}
        #chatBox {{ 
            position: absolute; bottom: 10px; left: 10px; width: 300px; height: 150px; 
            background: rgba(0, 0, 0, 0.7); color: white; padding: 10px; 
            border-radius: 5px; overflow-y: auto; 
        }}
        #status {{ 
            position: absolute; top: 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="status">Players: 1</div>
        <div id="chatBox"></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);
        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;
        scene.add(ground);

        let players = {{}};
        const playerMeshes = {{}};

        function updatePlayers(playerData) {{
            playerData.forEach(player => {{
                if (!playerMeshes[player.username]) {{
                    const geometry = new THREE.BoxGeometry(2, 2, 2);
                    const material = new THREE.MeshPhongMaterial({{ color: player.color }});
                    const mesh = new THREE.Mesh(geometry, material);
                    scene.add(mesh);
                    playerMeshes[player.username] = mesh;
                }}
                const mesh = playerMeshes[player.username];
                mesh.position.set(player.x, 1, player.z);
            }});
            document.getElementById('status').textContent = `Players: ${{Object.keys(playerMeshes).length}}`;
        }}

        document.addEventListener('keydown', (event) => {{
            const speed = 2;
            switch (event.code) {{
                case 'ArrowLeft': case 'KeyA': 
                    st.session_state.move_left = true; 
                    st.session_state.x_pos -= speed; 
                    ws.send(`{playerName}|MOVE:${{st.session_state.x_pos}}:${{st.session_state.z_pos}}`);
                    break;
                case 'ArrowRight': case 'KeyD': 
                    st.session_state.move_right = true; 
                    st.session_state.x_pos += speed; 
                    ws.send(`{playerName}|MOVE:${{st.session_state.x_pos}}:${{st.session_state.z_pos}}`);
                    break;
                case 'ArrowUp': case 'KeyW': 
                    st.session_state.move_up = true; 
                    st.session_state.z_pos -= speed; 
                    ws.send(`{playerName}|MOVE:${{st.session_state.x_pos}}:${{st.session_state.z_pos}}`);
                    break;
                case 'ArrowDown': case 'KeyS': 
                    st.session_state.move_down = true; 
                    st.session_state.z_pos += speed; 
                    ws.send(`{playerName}|MOVE:${{st.session_state.x_pos}}:${{st.session_state.z_pos}}`);
                    break;
            }}
        }});

        document.addEventListener('keyup', (event) => {{
            switch (event.code) {{
                case 'ArrowLeft': case 'KeyA': st.session_state.move_left = false; break;
                case 'ArrowRight': case 'KeyD': st.session_state.move_right = false; break;
                case 'ArrowUp': case 'KeyW': st.session_state.move_up = false; break;
                case 'ArrowDown': case 'KeyS': st.session_state.move_down = false; break;
            }}
        }});

        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 {{
                const [sender, message] = data.split('|');
                const chatBox = document.getElementById('chatBox');
                chatBox.innerHTML += `<p>${{sender}}: ${{message}}</p>`;
                chatBox.scrollTop = chatBox.scrollHeight;
            }}
        }};

        function animate() {{
            requestAnimationFrame(animate);
            renderer.render(scene, camera);
        }}
        animate();
    </script>
</body>
</html>
"""

# Main Game Loop
def main():
    st.title(f"🎮 {GAME_NAME}")
    st.subheader(f"🌄 {st.session_state.username}’s Adventure - Score: {st.session_state.score} 🏆")

    # Countdown Timer
    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")
    if remaining <= 0:
        st.session_state.last_update = time.time()
        st.rerun()

    # Voice Input Component
    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 = CHARACTERS.get(st.session_state.username, {"voice": "en-US-AriaNeural"})["voice"]
            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.rerun()

    # Render Map
    components.html(html_code, width=800, height=600)

    # Chat Interface
    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 = CHARACTERS[st.session_state.username]["voice"]
            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.success(f"🌄 +10 points! New Score: {st.session_state.score}")

    # Sidebar: Game HUD
    st.sidebar.subheader("🎮 Adventurer’s HUD")
    new_username = st.sidebar.selectbox("🧙‍♂️ Choose Your Hero", list(CHARACTERS.keys()), index=list(CHARACTERS.keys()).index(st.session_state.username))
    if new_username != st.session_state.username:
        asyncio.run(save_chat_entry(st.session_state.username, f"🔄 Transforms into {new_username}!", CHARACTERS[st.session_state.username]["voice"]))
        st.session_state.username = new_username
        save_username(st.session_state.username)
        st.rerun()
    st.sidebar.write(f"📜 {CHARACTERS[st.session_state.username]['desc']}")
    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'}")

    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()