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  # For regular expressions in clean_text_for_tts
import threading  # For running WebSocket server in a separate thread
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!"},
    "Meme Queen Mia 😂": {"voice": "en-US-JennyNeural", "desc": "Spreads laughs with wild memes!"},
    "Elk Whisperer Eve 🦌": {"voice": "en-GB-SoniaNeural", "desc": "Talks to wildlife, loves nature!"},
    "Tech Titan Tara 💾": {"voice": "en-AU-NatashaNeural", "desc": "Codes her way through the Rockies!"},
    "Ski Guru Sam ⛷️": {"voice": "en-CA-ClaraNeural", "desc": "Shreds slopes, lives for snow!"},
    "Cosmic Camper Cal 🌠": {"voice": "en-US-AriaNeural", "desc": "Stargazes and tells epic tales!"},
    "Rasta Ranger Rick 🍃": {"voice": "en-GB-RyanNeural", "desc": "Chills with nature’s vibes!"},
    "Boulder Bro Ben 🪨": {"voice": "en-AU-WilliamNeural", "desc": "Climbs rocks, bro-style!"}
}
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")

# 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,
        'marquee_settings': {
            "background": "#2E8B57", "color": "#FFFFFF", "font-size": "16px",
            "animationDuration": "15s", "width": "100%", "lineHeight": "40px"
        }
    }
    for k, v in defaults.items():
        if k not in st.session_state:
            st.session_state[k] = v

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

# 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  # Points for participation
    st.session_state.treasures += 1  # Audio treasure collected
    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')

# 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
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.get('username', random.choice(list(CHARACTERS.keys())))
    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"]
                await save_chat_entry(username, content, voice)
                await perform_arxiv_search(content, username)  # ArXiv response for every chat
    except websockets.ConnectionClosed:
        await save_chat_entry(username, "🏃 Leaves the quest!", CHARACTERS[username]["voice"])
    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 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
        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 Quest (Mad Libs)
def generate_quest_story(username, inputs):
    locations = ["Peak Summit 🏞️", "Elk Valley 🌲", "Meme Cave 🕳️", "Tech Outpost 💻"]
    st.session_state.location = random.choice(locations)
    story = f"""
    🌄 **{username}’s Quest at {st.session_state.location}:**
    🗺️ {username} discovers {inputs['quantity']} {inputs['plural_noun']}!
    🌟 They’re {inputs['adjective']} and {inputs['action']} everywhere.
    🎯 With a {inputs['tool']}, {username} faces a {inputs['creature']} 
    ⚡ that {inputs['event']}—shouting “{inputs['shout']}!”
    🎉 Victory earns {username} {st.session_state.treasures} audio treasures!
    """
    return story.strip()

# Main Game Loop
def main():
    init_session_state()
    saved_username = load_username()
    if saved_username and saved_username in CHARACTERS:
        st.session_state.username = saved_username
    if not st.session_state.username:
        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)

    st.title(f"🎮 {GAME_NAME}")
    st.subheader(f"🌄 {st.session_state.username}’s Adventure - Score: {st.session_state.score} 🏆")
    chat_text = " ".join([line.split(": ")[-1] for line in asyncio.run(load_chat()) if ": " in line][-5:])
    streamlit_marquee(content=f"🏔️ {st.session_state.location} | 🎙️ {st.session_state.username} | 💬 {chat_text}", 
                      **st.session_state['marquee_settings'], key="quest_marquee")

    tab_main = st.radio("🎲 Quest Actions:", ["🗣️ Explore & Chat", "🎵 Treasure Vault", "🗺️ Quest Challenge"], horizontal=True)

    if tab_main == "🗣️ Explore & Chat":
        st.subheader(f"🗣️ Explore {st.session_state.location} 🎙️")
        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="Explore the wilds! 🌲")
        if st.button("🌟 Send & Explore 🎤"):
            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}")

    elif tab_main == "🎵 Treasure Vault":
        st.subheader("🎵 Audio Treasure Vault 🏆")
        mp3_files = sorted(glob.glob("*.mp3"), key=os.path.getmtime, reverse=True)
        if mp3_files:
            st.write(f"🏅 Treasures Collected: {st.session_state.treasures}")
            for i, mp3 in enumerate(mp3_files[:10]):
                with st.expander(f"🎵 Treasure #{i+1}: {os.path.basename(mp3)}"):
                    play_and_download_audio(mp3)
        else:
            st.write("🔍 No treasures yet—explore or complete quests to collect audio loot! 🎤")

    elif tab_main == "🗺️ Quest Challenge":
        st.subheader("🗺️ Rocky Mountain Quest Challenge 🌟")
        st.write("Fill in the blanks to embark on a wild adventure!")
        inputs = {}
        col1, col2 = st.columns(2)
        with col1:
            inputs['quantity'] = st.text_input("🔢 How Many?", "3", key="quantity")
            inputs['plural_noun'] = st.text_input("🌄 Things?", "elk", key="plural_noun")
            inputs['adjective'] = st.text_input("✨ Describe Them?", "wild", key="adjective")
            inputs['action'] = st.text_input("🏃 What They Do?", "running", key="action")
        with col2:
            inputs['tool'] = st.text_input("🛠️ Your Tool?", "map", key="tool")
            inputs['creature'] = st.text_input("🦇 Encounter?", "bear", key="creature")
            inputs['event'] = st.text_input("⚡ What Happens?", "roars", key="event")
            inputs['shout'] = st.text_input("🗣️ Your Cry?", "Yeehaw!", key="shout")

        if st.button("🎉 Start Quest! 🎤"):
            story = generate_quest_story(st.session_state.username, inputs)
            st.markdown(f"### 🌄 {st.session_state.username}’s Quest Log")
            st.write(story)
            voice = CHARACTERS[st.session_state.username]["voice"]
            md_file, audio_file = asyncio.run(save_chat_entry(st.session_state.username, story, voice, True))
            if audio_file:
                play_and_download_audio(audio_file)
            st.session_state.score += 50  # Bonus for quest completion
            st.success(f"🏆 Quest Complete! +50 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}")

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