import streamlit as st
from gradio_client import Client
import tempfile
import os
import time
import shutil

# Page configuration
st.set_page_config(
    page_title="VibeVoice TTS Tester",
    page_icon="🎤",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for modern styling
st.markdown("""
<style>
    .main-header {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        padding: 2rem;
        border-radius: 15px;
        margin-bottom: 2rem;
        text-align: center;
        color: white;
    }
    .main-header h1 {
        margin: 0;
        font-size: 2.5rem;
    }
    .main-header p {
        margin: 0.5rem 0 0 0;
        opacity: 0.9;
    }
    .main-header a {
        color: #ffd700;
        text-decoration: none;
        font-weight: bold;
    }
    .main-header a:hover {
        text-decoration: underline;
    }
    .stButton > button {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        border: none;
        padding: 0.75rem 2rem;
        font-size: 1.1rem;
        border-radius: 10px;
        transition: all 0.3s ease;
        width: 100%;
    }
    .stButton > button:hover {
        transform: translateY(-2px);
        box-shadow: 0 5px 20px rgba(102, 126, 234, 0.4);
    }
    .info-box {
        background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
        padding: 1.5rem;
        border-radius: 10px;
        margin: 1rem 0;
    }
    .audio-container {
        background: #f8f9fa;
        padding: 1.5rem;
        border-radius: 10px;
        margin: 1rem 0;
        border: 2px solid #667eea;
    }
    .status-success {
        color: #28a745;
        font-weight: bold;
    }
    .status-error {
        color: #dc3545;
        font-weight: bold;
    }
    .example-text {
        background: #f0f2f6;
        padding: 0.5rem 1rem;
        border-radius: 5px;
        margin: 0.25rem 0;
        cursor: pointer;
        transition: all 0.2s ease;
    }
    .example-text:hover {
        background: #e0e2e6;
    }
</style>
""", unsafe_allow_html=True)

# Header
st.markdown("""
<div class="main-header">
    <h1>🎤 VibeVoice TTS Tester</h1>
    <p>Test the VibeVoice-Realtime-0.5B Text-to-Speech Model</p>
    <p style="margin-top: 1rem; font-size: 0.9rem;">
        <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank">Built with anycoder</a>
    </p>
</div>
""", unsafe_allow_html=True)

# Initialize session state
if 'audio_history' not in st.session_state:
    st.session_state.audio_history = []
if 'client' not in st.session_state:
    st.session_state.client = None
if 'connection_status' not in st.session_state:
    st.session_state.connection_status = None
if 'current_text' not in st.session_state:
    st.session_state.current_text = ""

# Sidebar
with st.sidebar:
    st.markdown("### ⚙️ Settings")
    
    # Connection status
    st.markdown("---")
    st.markdown("### 🔌 Connection Status")
    
    if st.button("🔄 Test Connection", use_container_width=True):
        with st.spinner("Connecting to VibeVoice API..."):
            try:
                client = Client("anycoderapps/VibeVoice-Realtime-0.5B")
                st.session_state.client = client
                st.session_state.connection_status = "connected"
                st.success("✅ Connected successfully!")
            except Exception as e:
                st.session_state.connection_status = "error"
                st.error(f"❌ Connection failed: {str(e)}")
    
    if st.session_state.connection_status == "connected":
        st.markdown('<p class="status-success">● Connected</p>', unsafe_allow_html=True)
    elif st.session_state.connection_status == "error":
        st.markdown('<p class="status-error">● Disconnected</p>', unsafe_allow_html=True)
    else:
        st.markdown('<p style="color: #6c757d;">● Not tested</p>', unsafe_allow_html=True)
    
    st.markdown("---")
    st.markdown("### 📚 About")
    st.markdown("""
    This app allows you to test the **VibeVoice-Realtime-0.5B** 
    text-to-speech model hosted on Hugging Face Spaces.
    
    **Features:**
    - 🎯 Real-time TTS generation
    - 🎵 Audio playback & download
    - 📝 Example prompts
    - 📊 Generation history
    """)
    
    st.markdown("---")
    st.markdown("### 🔗 Links")
    st.markdown("[VibeVoice Space](https://huggingface.co/spaces/anycoderapps/VibeVoice-Realtime-0.5B)")
    
    # Clear history button
    st.markdown("---")
    if st.button("🗑️ Clear History", use_container_width=True):
        st.session_state.audio_history = []
        st.rerun()

# Main content
col1, col2 = st.columns([2, 1])

with col1:
    st.markdown("### 📝 Enter Your Text")
    
    # Text input
    text_input = st.text_area(
        "Text to synthesize",
        value=st.session_state.current_text,
        placeholder="Enter the text you want to convert to speech...",
        height=150,
        help="Enter any text and the TTS model will convert it to natural-sounding speech."
    )
    
    # Character count
    char_count = len(text_input)
    st.caption(f"Characters: {char_count}")
    
    # Generate button
    generate_button = st.button("🎤 Generate Speech", use_container_width=True, type="primary")

with col2:
    st.markdown("### 💡 Example Prompts")
    
    example_texts = [
        "Hello! Welcome to VibeVoice, the next generation text-to-speech system.",
        "The quick brown fox jumps over the lazy dog.",
        "Artificial intelligence is transforming the way we interact with technology.",
        "Good morning! I hope you have a wonderful day ahead.",
        "This is a test of the emergency broadcast system.",
    ]
    
    for i, example in enumerate(example_texts):
        if st.button(f"📌 Example {i+1}", key=f"example_{i}", use_container_width=True):
            st.session_state.current_text = example
            st.rerun()

# Generate speech
if generate_button:
    if not text_input.strip():
        st.error("⚠️ Please enter some text to synthesize.")
    else:
        with st.spinner("🎵 Generating speech... This may take a moment."):
            try:
                # Initialize client if not already done
                if st.session_state.client is None:
                    st.session_state.client = Client("anycoderapps/VibeVoice-Realtime-0.5B")
                    st.session_state.connection_status = "connected"
                
                client = st.session_state.client
                
                # Call the API
                start_time = time.time()
                result = client.predict(
                    text=text_input,
                    api_name="/predict"
                )
                end_time = time.time()
                
                generation_time = round(end_time - start_time, 2)
                
                # Handle the result
                if result:
                    st.success(f"✅ Speech generated successfully in {generation_time}s!")
                    
                    # Read audio bytes for storage
                    audio_bytes = None
                    if os.path.exists(result):
                        with open(result, 'rb') as audio_file:
                            audio_bytes = audio_file.read()
                    
                    # Store in history with audio bytes
                    st.session_state.audio_history.insert(0, {
                        'text': text_input[:100] + "..." if len(text_input) > 100 else text_input,
                        'full_text': text_input,
                        'audio_bytes': audio_bytes,
                        'timestamp': time.strftime("%H:%M:%S"),
                        'generation_time': generation_time
                    })
                    
                    # Keep only last 10 items
                    st.session_state.audio_history = st.session_state.audio_history[:10]
                    
                    # Display audio
                    st.markdown("### 🎵 Generated Audio")
                    
                    if audio_bytes:
                        st.audio(audio_bytes, format='audio/wav')
                        
                        # Download button
                        st.download_button(
                            label="⬇️ Download Audio",
                            data=audio_bytes,
                            file_name="vibevoice_output.wav",
                            mime="audio/wav",
                            use_container_width=True
                        )
                    else:
                        st.warning("Audio file path returned but file not accessible.")
                        st.info(f"Result path: {result}")
                else:
                    st.error("❌ No audio was generated. Please try again.")
                    
            except Exception as e:
                st.error(f"❌ Error generating speech: {str(e)}")
                st.session_state.connection_status = "error"
                st.session_state.client = None  # Reset client on error
                
                # Show troubleshooting tips
                with st.expander("🔧 Troubleshooting Tips"):
                    st.markdown("""
                    1. **Check the Space**: Make sure the [VibeVoice Space](https://huggingface.co/spaces/anycoderapps/VibeVoice-Realtime-0.5B) is running.
                    2. **Wait and Retry**: The Space might be starting up. Wait a minute and try again.
                    3. **Shorter Text**: Try with a shorter text input.
                    4. **Refresh Connection**: Click "Test Connection" in the sidebar.
                    """)

# Display history
st.markdown("---")
st.markdown("### 📊 Generation History")

if st.session_state.audio_history:
    for i, item in enumerate(st.session_state.audio_history):
        with st.expander(f"🕐 {item['timestamp']} - {item['text']}", expanded=(i == 0)):
            col_a, col_b = st.columns([3, 1])
            with col_a:
                st.markdown(f"**Text:** {item.get('full_text', item['text'])}")
                st.caption(f"Generation time: {item['generation_time']}s")
            with col_b:
                if item.get('audio_bytes'):
                    st.audio(item['audio_bytes'], format='audio/wav')
                    st.download_button(
                        label="⬇️",
                        data=item['audio_bytes'],
                        file_name=f"vibevoice_{item['timestamp'].replace(':', '-')}.wav",
                        mime="audio/wav",
                        key=f"download_{i}",
                        use_container_width=True
                    )
else:
    st.info("No audio generated yet. Enter some text and click 'Generate Speech' to get started!")

# Footer
st.markdown("---")
st.markdown("""
<div style="text-align: center; color: #6c757d; padding: 1rem;">
    <p>Powered by <a href="https://huggingface.co/spaces/anycoderapps/VibeVoice-Realtime-0.5B" target="_blank">VibeVoice-Realtime-0.5B</a></p>
    <p style="font-size: 0.9rem;">
        <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank">Built with anycoder</a>
    </p>
</div>
""", unsafe_allow_html=True)