import streamlit as st
import streamlit.components.v1 as components
import asyncio
import edge_tts
import os
import base64
import json
from datetime import datetime
from typing import Optional, Dict, List
import glob

# Configure page
st.set_page_config(
    page_title="Research Assistant",
    page_icon="🔬",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Initialize session state
if 'current_query' not in st.session_state:
    st.session_state.current_query = ''
if 'last_response' not in st.session_state:
    st.session_state.last_response = ''
if 'audio_queue' not in st.session_state:
    st.session_state.audio_queue = []
if 'mermaid_history' not in st.session_state:
    st.session_state.mermaid_history = []

# Custom CSS
st.markdown("""
<style>
    .main { background-color: #f5f5f5; }
    .stMarkdown { font-family: 'Inter', sans-serif; }
    .diagram-container { 
        background: white;
        padding: 20px;
        border-radius: 10px;
        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
    }
    .results-container {
        margin-top: 20px;
        padding: 15px;
        background: white;
        border-radius: 10px;
    }
    .audio-player {
        width: 100%;
        margin: 10px 0;
    }
</style>
""", unsafe_allow_html=True)

def generate_mermaid_html(mermaid_code: str, height: int = 400) -> str:
    """Generate responsive Mermaid diagram HTML with click handling."""
    return f"""
    <div class="diagram-container">
        <script src="https://cdn.jsdelivr.net/npm/mermaid/dist/mermaid.min.js"></script>
        <div class="mermaid" style="height: {height}px;">
            {mermaid_code}
        </div>
        <script>
            mermaid.initialize({{
                startOnLoad: true,
                securityLevel: 'loose',
                theme: 'default'
            }});
            
            document.addEventListener('click', (e) => {{
                if (e.target.tagName === 'g' && e.target.classList.contains('node')) {{
                    const nodeId = e.target.id;
                    window.parent.postMessage({{
                        type: 'node_clicked',
                        nodeId: nodeId,
                        isStreamlitMessage: true
                    }}, '*');
                }}
            }});
        </script>
    </div>
    """

async def generate_speech(text: str, voice: str = "en-US-AriaNeural") -> Optional[str]:
    """Generate speech using Edge TTS."""
    if not text.strip():
        return None
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_file = f"speech_{timestamp}.mp3"
    
    communicate = edge_tts.Communicate(text, voice)
    await communicate.save(output_file)
    
    return output_file

def process_arxiv_search(query: str) -> Dict:
    """Process Arxiv search with your existing research code."""
    # Integrate your Arxiv search code here
    # This is a placeholder that simulates a response
    return {
        "title": "Sample Research Paper",
        "abstract": "This is a sample abstract...",
        "authors": ["Author 1", "Author 2"],
        "url": "https://arxiv.org/abs/..."
    }

def create_audio_player(file_path: str) -> str:
    """Create an HTML audio player with download button."""
    with open(file_path, "rb") as audio_file:
        audio_bytes = audio_file.read()
    audio_b64 = base64.b64encode(audio_bytes).decode()
    
    return f"""
    <div class="audio-player">
        <audio controls style="width: 100%">
            <source src="data:audio/mp3;base64,{audio_b64}" type="audio/mp3">
            Your browser does not support the audio element.
        </audio>
        <a href="data:audio/mp3;base64,{audio_b64}" 
           download="{os.path.basename(file_path)}"
           style="margin-top: 5px; display: inline-block;">
           Download Audio
        </a>
    </div>
    """

def handle_node_click(node_id: str):
    """Handle Mermaid diagram node clicks."""
    # Convert node ID to search query
    query = node_id.replace('_', ' ')
    
    # Perform search
    results = process_arxiv_search(query)
    
    # Generate speech from results
    asyncio.run(generate_speech(results['abstract']))
    
    # Update session state
    st.session_state.current_query = query
    st.session_state.last_response = results

# Main Mermaid diagram definition
RESEARCH_DIAGRAM = """
graph TD
    A[Literature Review] --> B[Data Analysis]
    B --> C[Results]
    C --> D[Conclusions]
    
    click A callback "Research Methodology"
    click B callback "Statistical Analysis"
    click C callback "Research Findings"
    click D callback "Research Impact"
    
    style A fill:#f9f,stroke:#333,stroke-width:4px
    style B fill:#bbf,stroke:#333,stroke-width:4px
    style C fill:#bfb,stroke:#333,stroke-width:4px
    style D fill:#fbb,stroke:#333,stroke-width:4px
"""

def main():
    st.title("📚 Research Assistant")
    
    # Sidebar configuration
    st.sidebar.header("Configuration")
    voice_option = st.sidebar.selectbox(
        "Select Voice",
        ["en-US-AriaNeural", "en-US-GuyNeural", "en-GB-SoniaNeural"]
    )
    
    # Main layout
    col1, col2 = st.columns([2, 3])
    
    with col1:
        st.subheader("Research Map")
        components.html(
            generate_mermaid_html(RESEARCH_DIAGRAM),
            height=500,
            scrolling=True
        )
        
        st.markdown("### Recent Searches")
        for query in st.session_state.mermaid_history[-5:]:
            st.info(query)
    
    with col2:
        st.subheader("Research Results")
        
        # Manual search option
        search_query = st.text_input("Enter search query:")
        if st.button("Search"):
            handle_node_click(search_query)
        
        # Display current results
        if st.session_state.last_response:
            with st.container():
                st.markdown("#### Latest Results")
                st.json(st.session_state.last_response)
                
                # Audio playback
                audio_files = glob.glob("speech_*.mp3")
                if audio_files:
                    latest_audio = max(audio_files, key=os.path.getctime)
                    st.markdown(create_audio_player(latest_audio), unsafe_allow_html=True)
    
    # Cleanup old audio files
    for file in glob.glob("speech_*.mp3")[:-5]:  # Keep only last 5 files
        try:
            os.remove(file)
        except:
            pass

if __name__ == "__main__":
    main()