import streamlit as st
import asyncio
import websockets
import uuid
import argparse
from datetime import datetime
import os
import random
import time
import hashlib
from PIL import Image
import glob
import base64
import io
import streamlit.components.v1 as components
import edge_tts
from audio_recorder_streamlit import audio_recorder
import nest_asyncio
import re
from streamlit_paste_button import paste_image_button
import pytz
import shutil
import logging

# Set up basic logging
logging.basicConfig(level=logging.DEBUG, filename="app.log", filemode="a",
                    format="%(asctime)s - %(levelname)s - %(message)s")

# Patch for nested async
nest_asyncio.apply()

# Static config
icons = '🤖🧠🔬📝'
START_ROOM = "Sector 🌌"

# Page setup
st.set_page_config(
    page_title="🤖🧠MMO Chat Brain📝🔬",
    page_icon=icons,
    layout="wide",
    initial_sidebar_state="auto"
)

# Funky usernames with 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",
    "ByteBuster 💾": "en-GB-RyanNeural",
    "GalacticGopher 🌍": "en-AU-WilliamNeural",
    "RocketRaccoon 🚀": "en-CA-LiamNeural",
    "EchoElf 🧝": "en-US-AnaNeural",
    "PhantomFox 🦊": "en-US-BrandonNeural",
    "WittyWizard 🧙": "en-GB-ThomasNeural",
    "LunarLlama 🌙": "en-AU-FreyaNeural",
    "SolarSloth ☀️": "en-CA-LindaNeural",
    "AstroAlpaca 🦙": "en-US-ChristopherNeural",
    "CyberCoyote 🐺": "en-GB-ElliotNeural",
    "MysticMoose 🦌": "en-AU-JamesNeural",
    "GlitchGnome 🧚": "en-CA-EthanNeural",
    "VortexViper 🐍": "en-US-AmberNeural",
    "ChronoChimp 🐒": "en-GB-LibbyNeural"
}

# Folders
CHAT_DIR = "chat_logs"
VOTE_DIR = "vote_logs"
STATE_FILE = "user_state.txt"
AUDIO_DIR = "audio_logs"
HISTORY_DIR = "history_logs"
MEDIA_DIR = "media_files"
for dir in [CHAT_DIR, VOTE_DIR, AUDIO_DIR, HISTORY_DIR, MEDIA_DIR]:
    os.makedirs(dir, exist_ok=True)

CHAT_FILE = os.path.join(CHAT_DIR, "global_chat.md")
QUOTE_VOTES_FILE = os.path.join(VOTE_DIR, "quote_votes.md")
MEDIA_VOTES_FILE = os.path.join(VOTE_DIR, "media_votes.md")
HISTORY_FILE = os.path.join(HISTORY_DIR, "chat_history.md")

# Unicode digits
UNICODE_DIGITS = {i: f"{i}\uFE0F⃣" for i in range(10)}

# Font collection (simplified for brevity)
UNICODE_FONTS = [
    ("Normal", lambda x: x),
    ("Bold", lambda x: "".join(chr(ord(c) + 0x1D400 - 0x41) if 'A' <= c <= 'Z' else chr(ord(c) + 0x1D41A - 0x61) if 'a' <= c <= 'z' else c for c in x)),
]

# Global state
if 'server_running' not in st.session_state:
    st.session_state.server_running = False
if 'server_task' not in st.session_state:
    st.session_state.server_task = None
if 'active_connections' not in st.session_state:
    st.session_state.active_connections = {}
if 'media_notifications' not in st.session_state:
    st.session_state.media_notifications = []
if 'last_chat_update' not in st.session_state:
    st.session_state.last_chat_update = 0
if 'displayed_chat_lines' not in st.session_state:
    st.session_state.displayed_chat_lines = []
if 'message_text' not in st.session_state:
    st.session_state.message_text = ""
if 'audio_cache' not in st.session_state:
    st.session_state.audio_cache = {}
if 'pasted_image_data' not in st.session_state:
    st.session_state.pasted_image_data = None
if 'quote_line' not in st.session_state:
    st.session_state.quote_line = None
if 'refresh_rate' not in st.session_state:
    st.session_state.refresh_rate = 5
if 'base64_cache' not in st.session_state:
    st.session_state.base64_cache = {}
if 'transcript_history' not in st.session_state:
    st.session_state.transcript_history = []
if 'last_transcript' not in st.session_state:
    st.session_state.last_transcript = ""
if 'image_hashes' not in st.session_state:
    st.session_state.image_hashes = set()

# Utility functions
def format_timestamp_prefix(username):
    central = pytz.timezone('US/Central')
    now = datetime.now(central)
    return f"{now.strftime('%I-%M-%p-ct-%m-%d-%Y')}-by-{username}"

def compute_image_hash(image_data):
    if isinstance(image_data, Image.Image):
        img_byte_arr = io.BytesIO()
        image_data.save(img_byte_arr, format='PNG')
        img_bytes = img_byte_arr.getvalue()
    else:
        img_bytes = image_data
    return hashlib.md5(img_bytes).hexdigest()[:8]

def get_node_name():
    parser = argparse.ArgumentParser(description='Start a chat node with a specific name')
    parser.add_argument('--node-name', type=str, default=None)
    parser.add_argument('--port', type=int, default=8501)
    args = parser.parse_args()
    return args.node_name or f"node-{uuid.uuid4().hex[:8]}", args.port

def log_action(username, action):
    logging.debug(f"{username}: {action}")
    with open(HISTORY_FILE, 'a') as f:
        central = pytz.timezone('US/Central')
        f.write(f"[{datetime.now(central).strftime('%Y-%m-%d %H:%M:%S')}] {username}: {action}\n")

def clean_text_for_tts(text):
    cleaned = re.sub(r'[#*!\[\]]+', '', text)
    cleaned = ' '.join(cleaned.split())
    return cleaned[:200] if cleaned else "No text to speak"

async def save_chat_entry(username, message, is_markdown=False):
    try:
        log_action(username, "Saving chat entry")
        central = pytz.timezone('US/Central')
        timestamp = datetime.now(central).strftime("%Y-%m-%d %H:%M:%S")
        entry = f"[{timestamp}] {username}:\n```markdown\n{message}\n```" if is_markdown else f"[{timestamp}] {username}: {message}"
        with open(CHAT_FILE, 'a') as f:
            f.write(f"{entry}\n")
        voice = FUN_USERNAMES.get(username, "en-US-AriaNeural")
        cleaned_message = clean_text_for_tts(message)
        audio_file = await async_edge_tts_generate(cleaned_message, voice)
        if audio_file:
            with open(HISTORY_FILE, 'a') as f:
                f.write(f"[{timestamp}] {username}: Audio generated - {audio_file}\n")
        await broadcast_message(f"{username}|{message}", "chat")
        st.session_state.last_chat_update = time.time()
        return audio_file
    except Exception as e:
        logging.error(f"Error in save_chat_entry: {str(e)}")
        return None

async def load_chat():
    try:
        if not os.path.exists(CHAT_FILE):
            with open(CHAT_FILE, 'a') as f:
                f.write(f"# {START_ROOM} Chat\n\nWelcome to the cosmic hub - start chatting! 🎤\n")
        with open(CHAT_FILE, 'r') as f:
            return f.read()
    except Exception as e:
        logging.error(f"Error in load_chat: {str(e)}")
        return ""

async def save_audio_recording(transcript, username, voice):
    try:
        timestamp = format_timestamp_prefix(username)
        voice_id = voice.split('-')[-1].lower()
        filename = f"rec_{username}_{voice_id}_{timestamp}.mp3"
        filepath = os.path.join(AUDIO_DIR, filename)
        
        # Use edge_tts to generate audio from transcript since we don't have raw audio
        audio_file = await async_edge_tts_generate(transcript or "Audio recording", voice, file_format="mp3")
        if audio_file and os.path.exists(audio_file):
            os.rename(audio_file, filepath)
            return filepath
        return None
    except Exception as e:
        logging.error(f"Error in save_audio_recording: {str(e)}")
        return None

async def async_edge_tts_generate(text, voice, rate=0, pitch=0, file_format="mp3"):
    try:
        timestamp = format_timestamp_prefix(st.session_state.get('username', 'System 🌟'))
        filename = f"{timestamp}.{file_format}"
        filepath = os.path.join(AUDIO_DIR, filename)
        communicate = edge_tts.Communicate(text, voice, rate=f"{rate:+d}%", pitch=f"{pitch:+d}Hz")
        await communicate.save(filepath)
        return filepath if os.path.exists(filepath) else None
    except Exception as e:
        logging.error(f"Error in async_edge_tts_generate: {str(e)}")
        return None

def play_and_download_audio(file_path):
    if file_path and os.path.exists(file_path):
        st.audio(file_path)
        if file_path not in st.session_state.base64_cache:
            with open(file_path, "rb") as f:
                b64 = base64.b64encode(f.read()).decode()
            st.session_state.base64_cache[file_path] = b64
        b64 = st.session_state.base64_cache[file_path]
        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)

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:
            del st.session_state.active_connections[room_id][client_id]

async def websocket_handler(websocket, path):
    try:
        client_id = str(uuid.uuid4())
        room_id = "chat"
        st.session_state.active_connections.setdefault(room_id, {})[client_id] = websocket
        chat_content = await load_chat()
        username = st.session_state.get('username', random.choice(list(FUN_USERNAMES.keys())))
        if not any(f"Client-{client_id}" in line for line in chat_content.split('\n')):
            await save_chat_entry(f"Client-{client_id}", f"{username} has joined {START_ROOM}!")
        async for message in websocket:
            parts = message.split('|', 1)
            if len(parts) == 2:
                username, content = parts
                await save_chat_entry(username, content)
    except Exception as e:
        logging.error(f"Error in websocket_handler: {str(e)}")
    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 run_websocket_server():
    try:
        if not st.session_state.server_running:
            server = await websockets.serve(websocket_handler, '0.0.0.0', 8765)
            st.session_state.server_running = True
            await server.wait_closed()
    except Exception as e:
        logging.error(f"Error in run_websocket_server: {str(e)}")

ASR_HTML = """
<html>
<head>
    <title>Continuous Speech Demo</title>
    <style>
        body { font-family: sans-serif; padding: 20px; max-width: 800px; margin: 0 auto; }
        button { padding: 10px 20px; margin: 10px 5px; font-size: 16px; }
        #status { margin: 10px 0; padding: 10px; background: #e8f5e9; border-radius: 4px; }
        #output { white-space: pre-wrap; padding: 15px; background: #f5f5f5; border-radius: 4px; margin: 10px 0; min-height: 100px; max-height: 400px; overflow-y: auto; }
        .controls { margin: 10px 0; }
    </style>
</head>
<body>
    <div class="controls">
        <button id="start">Start Listening</button>
        <button id="stop" disabled>Stop Listening</button>
        <button id="clear">Clear Text</button>
    </div>
    <div id="status">Ready</div>
    <div id="output"></div>
    
    <script>
        if (!('webkitSpeechRecognition' in window)) {
            alert('Speech recognition not supported');
        } else {
            const recognition = new webkitSpeechRecognition();
            const startButton = document.getElementById('start');
            const stopButton = document.getElementById('stop');
            const clearButton = document.getElementById('clear');
            const status = document.getElementById('status');
            const output = document.getElementById('output');
            let fullTranscript = '';
            let lastUpdateTime = Date.now();

            recognition.continuous = true;
            recognition.interimResults = true;

            const startRecognition = () => {
                try {
                    recognition.start();
                    status.textContent = 'Listening...';
                    startButton.disabled = true;
                    stopButton.disabled = false;
                } catch (e) {
                    console.error(e);
                    status.textContent = 'Error: ' + e.message;
                }
            };

            startButton.onclick = startRecognition;

            stopButton.onclick = () => {
                recognition.stop();
                status.textContent = 'Stopped';
                startButton.disabled = false;
                stopButton.disabled = true;
                if (fullTranscript) {
                    sendDataToPython({value: fullTranscript, dataType: "json", stopped: true});
                }
            };

            clearButton.onclick = () => {
                fullTranscript = '';
                output.textContent = '';
                sendDataToPython({value: '', dataType: "json"});
            };

            recognition.onresult = (event) => {
                let interimTranscript = '';
                let finalTranscript = '';

                for (let i = event.resultIndex; i < event.results.length; i++) {
                    const transcript = event.results[i][0].transcript;
                    if (event.results[i].isFinal) {
                        finalTranscript += transcript + '\\n';
                    } else {
                        interimTranscript += transcript;
                    }
                }

                if (finalTranscript || (Date.now() - lastUpdateTime > 5000)) {
                    if (finalTranscript) {
                        fullTranscript += finalTranscript;
                    }
                    lastUpdateTime = Date.now();
                    output.textContent = fullTranscript + (interimTranscript ? '... ' + interimTranscript : '');
                    output.scrollTop = output.scrollHeight;
                    sendDataToPython({value: fullTranscript, dataType: "json", stopped: false});
                }
            };

            recognition.onend = () => {
                if (!stopButton.disabled) {
                    try {
                        recognition.start();
                    } catch (e) {
                        console.error('Failed to restart recognition:', e);
                        status.textContent = 'Error restarting: ' + e.message;
                        startButton.disabled = false;
                        stopButton.disabled = true;
                    }
                }
            };

            recognition.onerror = (event) => {
                console.error('Recognition error:', event.error);
                status.textContent = 'Error: ' + event.error;
                startButton.disabled = false;
                stopButton.disabled = true;
            };
        }

        function sendDataToPython(data) {
            window.parent.postMessage({
                isStreamlitMessage: true,
                type: "streamlit:setComponentValue",
                ...data
            }, "*");
        }

        window.addEventListener('load', function() {
            window.setTimeout(function() {
                window.parent.postMessage({
                    isStreamlitMessage: true,
                    type: "streamlit:setFrameHeight",
                    height: document.documentElement.clientHeight
                }, "*");
            }, 0);
        });
    </script>
</body>
</html>
"""

def main():
    NODE_NAME, port = get_node_name()
    logging.info(f"Starting app with node name: {NODE_NAME}, port: {port}")

    if 'username' not in st.session_state:
        chat_content = asyncio.run(load_chat())
        available_names = [name for name in FUN_USERNAMES if not any(f"{name} has joined" in line for line in chat_content.split('\n'))]
        st.session_state.username = random.choice(available_names) if available_names else random.choice(list(FUN_USERNAMES.keys()))
        st.session_state.voice = FUN_USERNAMES[st.session_state.username]
        st.markdown(f"**🎙️ Voice Selected**: {st.session_state.voice} 🗣️ for {st.session_state.username}")

    st.title(f"🤖🧠MMO {st.session_state.username}📝🔬")
    st.markdown(f"Welcome to {START_ROOM} - chat, vote, upload, paste images, and enjoy quoting! 🎉")

    if not st.session_state.server_task:
        st.session_state.server_task = asyncio.create_task(run_websocket_server())

    audio_bytes = audio_recorder()
    if audio_bytes:
        audio_file = asyncio.run(process_voice_input(audio_bytes, st.session_state.username, st.session_state.voice))
        if audio_file:
            st.audio(audio_file)
            st.rerun()

    st.subheader("🎤 Continuous Speech Input")
    asr_component = components.html(ASR_HTML, height=400)
    if asr_component and isinstance(asr_component, dict) and 'value' in asr_component:
        transcript = asr_component['value'].strip()
        stopped = asr_component.get('stopped', False)
        if transcript and transcript != st.session_state.last_transcript:
            st.session_state.transcript_history.append(transcript)
            st.session_state.last_transcript = transcript
            if stopped:
                audio_file = asyncio.run(save_audio_recording(transcript, st.session_state.username, st.session_state.voice))
                asyncio.run(save_chat_entry(st.session_state.username, f"Voice message: {transcript}\nAudio file: {os.path.basename(audio_file)}", is_markdown=True))
                st.rerun()

    st.subheader("🎵 Recorded Audio Files")
    audio_files = glob.glob(f"{AUDIO_DIR}/rec_{st.session_state.username}_*.mp3")
    if audio_files:
        st.write(f"Found {len(audio_files)} recordings for {st.session_state.username}")
        for audio_file in sorted(audio_files, key=os.path.getmtime, reverse=True):
            col1, col2 = st.columns([3, 1])
            with col1:
                st.audio(audio_file)
                st.write(f"File: {os.path.basename(audio_file)}")
            with col2:
                play_and_download_audio(audio_file)

    st.subheader(f"{START_ROOM} Chat 💬")
    chat_content = asyncio.run(load_chat())
    chat_lines = chat_content.split('\n')
    for i, line in enumerate(chat_lines):
        if line.strip() and ': ' in line and not line.startswith('#'):
            st.markdown(line)

    message = st.text_input(f"Message as {st.session_state.username} (Voice: {st.session_state.voice})", key="message_input", value=st.session_state.message_text)
    if st.button("Send 🚀", key="send_button") and message.strip():
        asyncio.run(save_chat_entry(st.session_state.username, message, is_markdown=True))
        st.session_state.message_text = ''
        st.rerun()

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        logging.error(f"Main execution failed: {str(e)}")
        st.error(f"An error occurred: {str(e)}")