import os
import base64
import gradio as gr

# Import our custom modules
from src.audio import speech_to_text, text_to_speech
from src.rag import init_vector_store, retrieve_context_multi
from src.llm import generate_hindi_response, correct_hindi_query

# Initialize the RAG vector store on startup
print("Initializing application...")
init_vector_store("data")
print("Application ready.")


def process_voice_query(audio_filepath, history=None):
    """
    End-to-end pipeline: Audio In -> Audio Out + step-by-step text outputs.
    """
    if audio_filepath is None:
        return "🎙️ Please record your question first.", "", None

    # Step 1: Speech-to-Text
    hindi_transcript = speech_to_text(audio_filepath)
    if not hindi_transcript:
        return "Error: No speech detected. Please try again.", "", None
    if hindi_transcript.startswith("Error:"):
        return hindi_transcript, "", None

    # Step 2: Retrieve relevant FAQ context directly using the Hindi transcript
    # If the user says a short affirmation like "Haa" (Yes), use the previous bot topic for the RAG search
    search_query = hindi_transcript
    if history and len(hindi_transcript.split()) <= 3:
        last_user, last_bot, _ = history[-1]
        if last_bot and "क्या आप यह भी जानना चाहेंगे:" in last_bot:
            topic = last_bot.split("क्या आप यह भी जानना चाहेंगे:")[-1].strip().strip("?")
            if topic:
                search_query = topic
                print(f"[RAG] Short query detected. Using previous topic: {search_query}")

    # Auto-correct phonetic STT typos (e.g. 'laan' -> 'loan') before RAG search
    corrected_query = correct_hindi_query(search_query)
    print(f"[LLM] Corrected Query for RAG: {corrected_query}")

    faq_context = retrieve_context_multi("", corrected_query)

    # Step 3: Generate Hindi answer, using history for conversational context
    hindi_answer = generate_hindi_response(hindi_transcript, faq_context, history)

    # Step 4: Convert answer to spoken audio
    audio_output_path = text_to_speech(hindi_answer)

    return hindi_transcript, hindi_answer, audio_output_path


# ── Gradio UI (Mobile App Style) ──────────────────────────────────────────────
custom_css = """
/* Restrict main container to mobile width and center it */
.gradio-container {
    max-width: 450px !important;
    margin: auto !important;
    background-color: #f8fafc !important;
}

/* Card styling for elements */
.app-card {
    background: white;
    border-radius: 16px;
    padding: 16px;
    box-shadow: 0 4px 6px -1px rgb(0 0 0 / 0.1), 0 2px 4px -2px rgb(0 0 0 / 0.1);
    margin-bottom: 16px;
    border: 1px solid #e2e8f0;
}

/* Header styling */
.app-header {
    text-align: center;
    padding: 20px 0 10px 0;
    color: #0f172a;
}
.app-header h1 {
    font-size: 1.5rem;
    font-weight: 700;
    margin: 0;
}
.app-header p {
    color: #64748b;
    font-size: 0.9rem;
    margin-top: 4px;
}

/* Chat History Styling */
.chat-history {
    display: flex;
    flex-direction: column;
    gap: 16px;
    max-height: 400px;
    overflow-y: auto;
    padding: 4px;
}
.chat-turn {
    display: flex;
    flex-direction: column;
    gap: 6px;
}
.user-bubble {
    align-self: flex-end;
    background-color: #dcf8c6;
    color: #0f172a;
    padding: 10px 14px;
    border-radius: 16px 16px 0px 16px;
    max-width: 85%;
    box-shadow: 0 1px 2px rgb(0 0 0 / 0.1);
    font-size: 0.95rem;
}
.bot-row {
    display: flex;
    align-items: flex-start;
    gap: 8px;
    align-self: flex-start;
    max-width: 95%;
}
.bot-bubble {
    background-color: #f1f5f9;
    color: #0f172a;
    padding: 10px 14px;
    border-radius: 16px 16px 16px 0px;
    box-shadow: 0 1px 2px rgb(0 0 0 / 0.1);
    font-size: 0.95rem;
    flex: 1;
}
.play-btn {
    background: #059669;
    border: none;
    border-radius: 50%;
    width: 32px;
    height: 32px;
    cursor: pointer;
    display: flex;
    align-items: center;
    justify-content: center;
    flex-shrink: 0;
    margin-top: 4px;
    box-shadow: 0 1px 3px rgb(0 0 0 / 0.2);
}
.play-btn:hover { background: #047857; }
.play-btn:active { transform: scale(0.95); }
.chat-label {
    font-size: 0.75rem;
    color: #64748b;
    margin-bottom: 2px;
}

/* Hide footer */
footer { display: none !important; }
"""


def _audio_to_base64(audio_path):
    """Convert an audio file to a base64 data URI for inline playback."""
    if not audio_path or not os.path.exists(audio_path):
        return None
    try:
        with open(audio_path, "rb") as f:
            data = base64.b64encode(f.read()).decode("utf-8")
        return f"data:audio/mpeg;base64,{data}"
    except Exception:
        return None


def render_chat_history(history):
    """Render the full chat history as HTML with inline audio players."""
    if not history:
        return "<div class='chat-history'><div style='text-align:center; color:#94a3b8; font-size:0.9rem; margin-top:10px;'>Your conversation will appear here...</div></div>"

    html = "<div class='chat-history'>"
    for i, (user_msg, bot_msg, audio_b64) in enumerate(history):
        html += "<div class='chat-turn'>"

        # User bubble
        if user_msg:
            html += f"<div style='text-align:right'><div class='chat-label'>You 📝</div><div class='user-bubble'>{user_msg}</div></div>"

        # Bot bubble with speaker button
        if bot_msg:
            html += "<div><div class='chat-label'>Financial Assistant 📖</div><div class='bot-row'>"
            html += f"<div class='bot-bubble'>{bot_msg}</div>"
            if audio_b64:
                audio_id = f"audio_{i}"
                html += f"""
                    <button class='play-btn' onclick="
                        var a = document.getElementById('{audio_id}');
                        if (a.paused) {{ a.play(); this.innerHTML='⏸'; }}
                        else {{ a.pause(); this.innerHTML='🔊'; }}
                    " title="Play answer">🔊</button>
                    <audio id='{audio_id}' src='{audio_b64}'></audio>
                """
            html += "</div></div>"

        html += "</div>"
    html += "</div>"
    return html


def process_voice_query_ui(audio_filepath, history):
    """Process voice query and append to chat history."""
    if audio_filepath is None:
        return render_chat_history(history), None, history

    hindi_transcript, hindi_answer, audio_output_path = process_voice_query(audio_filepath, history)

    # Convert audio to base64 for inline playback in history
    audio_b64 = _audio_to_base64(audio_output_path)

    # Append to history
    history = history or []
    history.append((hindi_transcript, hindi_answer, audio_b64))

    chat_html = render_chat_history(history)
    return chat_html, audio_output_path, history


with gr.Blocks(title="Financial Assistant", css=custom_css, theme=gr.themes.Default(spacing_size="sm", radius_size="lg")) as demo:

    # Session state for chat history
    chat_history = gr.State([])

    # App Header
    gr.HTML(
        """
        <div class="app-header">
            <h1>🌾 Financial Assistant</h1>
            <p>Your simple banking helper</p>
        </div>
        """
    )

    # Top Card: Audio Controls
    with gr.Column(elem_classes="app-card"):
        gr.Markdown("### 🎙️ Ask a Question")
        audio_in = gr.Audio(
            sources=["microphone"],
            type="filepath",
            label="Tap to Record / Tap to Stop",
        )

        gr.Markdown("<br>### 🔊 Latest Answer")
        audio_out = gr.Audio(
            label="",
            interactive=False,
            autoplay=True,
        )

    # Bottom Card: Full Chat History
    with gr.Column(elem_classes="app-card"):
        gr.Markdown("### 💬 Conversation History")
        chat_display = gr.HTML(
            value=render_chat_history([])
        )

    # Trigger pipeline on recording stop
    audio_in.stop_recording(
        fn=process_voice_query_ui,
        inputs=[audio_in, chat_history],
        outputs=[chat_display, audio_out, chat_history],
    )

if __name__ == "__main__":
    demo.launch()