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
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 🏔️🎮",
    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": "🎵"}

# 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
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
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
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,
        'prairie_players': {}  # New for prairie simulator
    }
    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:
        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)

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
    
    # Handle both Rocky Mountain Quest and Prairie Simulator players
    if "prairie" in path:
        st.session_state.prairie_players[client_id] = {
            "username": username,
            "animal": "prairie_dog",  # Default
            "location": PRAIRIE_LOCATIONS["Deadwood, SD"],
            "color": CHARACTERS[username]["color"]
        }
    else:
        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)
                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"{username} ({st.session_state.prairie_players[client_id]['animal']}) moves to {value}"
                        await save_chat_entry(username, action_msg, voice)
                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]
        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"):
            # Rocky Mountain Quest update
            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")
            
            # Prairie Simulator update
            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"System|{prairie_message}", "quest")
            await broadcast_message(f"PRAIRIE_UPDATE:{prairie_data}", "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())

# Rocky Mountain Quest Map 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 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 = {{}};
        let xPos = 0;
        let zPos = 0;

        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':
                    xPos -= speed;
                    ws.send(`${{playerName}}|MOVE:${{xPos}}:${{zPos}}`);
                    break;
                case 'ArrowRight': case 'KeyD':
                    xPos += speed;
                    ws.send(`${{playerName}}|MOVE:${{xPos}}:${{zPos}}`);
                    break;
                case 'ArrowUp': case 'KeyW':
                    zPos -= speed;
                    ws.send(`${{playerName}}|MOVE:${{xPos}}:${{zPos}}`);
                    break;
                case 'ArrowDown': case 'KeyS':
                    zPos += speed;
                    ws.send(`${{playerName}}|MOVE:${{xPos}}:${{zPos}}`);
                    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 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;
            }}
        }};

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

# Prairie Simulator HTML
prairie_simulator_html = f"""
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Prairie Simulator</title>
    <style>
        body {{ margin: 0; font-family: Arial, sans-serif; background: #f0f0f0; }}
        #simContainer {{ width: 100%; height: 600px; position: relative; }}
        #map {{ width: 100%; height: 400px; }}
        #controls {{ padding: 10px; background: #fff; border-radius: 5px; }}
        #status {{ color: #333; padding: 5px; }}
        #chatBox {{ height: 100px; overflow-y: auto; background: #fff; border: 1px solid #ccc; padding: 5px; }}
    </style>
</head>
<body>
    <div id="simContainer">
        <div id="map"></div>
        <div id="controls">
            <label>Animal: </label>
            <select id="animal" onchange="updateAnimal()">
                <option value="prairie_dog">Prairie Dog</option>
                <option value="deer">Deer</option>
                <option value="sheep">Sheep</option>
                <option value="groundhog">Groundhog</option>
            </select>
            <label>Move to: </label>
            <select id="location">
                <option value="Deadwood, SD">Deadwood, SD</option>
                <option value="Wind Cave National Park">Wind Cave National Park</option>
                <option value="Wyoming Spring Creek">Wyoming Spring Creek</option>
            </select>
            <button onclick="move()">Move</button>
            <div id="status">Players: 0</div>
            <div id="chatBox"></div>
        </div>
    </div>

    <script src="https://unpkg.com/leaflet@1.9.4/dist/leaflet.js"></script>
    <link rel="stylesheet" href="https://unpkg.com/leaflet@1.9.4/dist/leaflet.css" />
    <script>
        const playerName = "{st.session_state.username}";
        let ws = new WebSocket('ws://localhost:8765/prairie');
        const map = L.map('map').setView([44.0, -103.0], 7);
        L.tileLayer('https://{s}.tile.openstreetmap.org/{z}/{x}/{y}.png', {{
            attribution: '&copy; OpenStreetMap contributors'
        }}).addTo(map);

        const locations = {{
            "Deadwood, SD": [44.3769, -103.7298],
            "Wind Cave National Park": [43.6047, -103.4798],
            "Wyoming Spring Creek": [41.6666, -106.6666]
        }};
        for (let loc in locations) {{
            L.marker(locations[loc]).addTo(map).bindPopup(loc);
        }}

        let players = {{}};
        const markers = {{}};

        function updatePlayers(playerData) {{
            playerData.forEach(player => {{
                if (!markers[player.username]) {{
                    markers[player.username] = L.circleMarker(player.location, {{
                        color: `#${{player.color.toString(16).padStart(6, '0')}}`,
                        radius: 5
                    }}).addTo(map).bindPopup(`${{player.username}} (${{player.animal}})`);
                }} else {{
                    markers[player.username].setLatLng(player.location);
                }}
            }});
            document.getElementById('status').textContent = `Players: ${{Object.keys(markers).length}}`;
        }}

        function updateAnimal() {{
            const animal = document.getElementById('animal').value;
            ws.send(`${{playerName}}|PRAIRIE:ANIMAL:${{animal}}`);
        }}

        function move() {{
            const location = document.getElementById('location').value;
            ws.send(`${{playerName}}|PRAIRIE:MOVE:${{location}}`);
        }}

        ws.onmessage = function(event) {{
            const data = event.data;
            if (data.startsWith('PRAIRIE_UPDATE:')) {{
                const playerData = JSON.parse(data.split('PRAIRIE_UPDATE:')[1]);
                updatePlayers(playerData);
            }} else if (!data.startsWith('MAP_UPDATE:')) {{
                const [sender, message] = data.split('|');
                const chatBox = document.getElementById('chatBox');
                chatBox.innerHTML += `<p>${{sender}}: ${{message}}</p>`;
                chatBox.scrollTop = chatBox.scrollHeight;
            }}
        }};
    </script>
</body>
</html>
"""

# Main Game Loop
def main():
    # Sidebar Titles and HUD
    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"📜 {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'}")
    
    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()

    # Two-column layout
    left_col, right_col = st.columns([2, 1])

    # Left Column: Rocky Mountain Quest
    with left_col:
        # Render Rocky Mountain Map
        components.html(rocky_map_html, 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}")

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

    # Right Column: Prairie Simulator
    with right_col:
        st.subheader("🌾 Prairie Simulator")
        components.html(prairie_simulator_html, width=600, height=600)

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

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