app.py

import gradio as gr
import os
import re
import tempfile
from datetime import datetime
from pathlib import Path
from dotenv import load_dotenv

from app.models import Message
from app.agent import get_agent

load_dotenv()

# ─── Groq TTS ───
GROQ_API_KEY = os.getenv("GROQ_API_KEY", "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")


def text_to_speech(history):
    """Convert the last assistant message to speech using Groq Orpheus TTS."""
    if not history:
        return None

    # Find the last assistant message
    last_bot_msg = None
    for turn in reversed(history):
        if turn.get("role") == "assistant":
            last_bot_msg = turn.get("content", "")
            break

    if not last_bot_msg:
        return None

    # Clean markdown for spoken text
    clean_text = strip_markdown(last_bot_msg)
    # Limit length for TTS (max ~500 chars to keep audio reasonable)
    if len(clean_text) > 500:
        clean_text = clean_text[:500] + "... That's the summary of my response."

    try:
        from groq import Groq

        client = Groq(api_key=GROQ_API_KEY)
        speech_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav", prefix="shl_tts_")

        response = client.audio.speech.create(
            model="canopylabs/orpheus-v1-english",
            voice="autumn",
            response_format="wav",
            input=clean_text,
        )
        response.write_to_file(speech_file.name)
        print(f"[TTS] Generated speech: {speech_file.name}")
        return speech_file.name

    except Exception as e:
        print(f"[TTS] Groq TTS failed: {e}. Falling back to gTTS.")
        try:
            from gtts import gTTS
            speech_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3", prefix="shl_tts_")
            tts = gTTS(text=clean_text, lang='en')
            tts.save(speech_file.name)
            print(f"[TTS] Generated fallback speech: {speech_file.name}")
            return speech_file.name
        except Exception as ge:
            print(f"[TTS] Fallback gTTS failed: {ge}")
            import traceback
            traceback.print_exc()
            return None


def transcribe_and_respond(audio_path, history):
    """Transcribe user audio input using Groq Whisper, then generate text & voice response."""
    if not audio_path:
        # Guard to prevent processing empty/cleared audio events
        return gr.update(), gr.update(), gr.update()
        
    try:
        from groq import Groq
        client = Groq(api_key=GROQ_API_KEY)
        
        print(f"[STT] Transcribing audio file: {audio_path}")
        with open(audio_path, "rb") as f:
            transcription = client.audio.transcriptions.create(
                file=(audio_path, f.read()),
                model="whisper-large-v3-turbo",
            )
        user_message = transcription.text
        print(f"[STT] Transcription result: '{user_message}'")
    except Exception as e:
        print(f"[STT] Transcription failed: {e}")
        import traceback
        traceback.print_exc()
        user_message = f"[Audio Transcription Error: {str(e)}]"

    # 1. Update history with user's transcribed query
    new_history = history + [{"role": "user", "content": user_message}]
    
    # 2. Generate bot response (using history before this turn)
    bot_response = respond(user_message, history)
    
    # 3. Update history with bot's response
    new_history.append({"role": "assistant", "content": bot_response})
    
    # 4. Generate TTS for the bot's response
    audio_path_out = text_to_speech(new_history)
    
    # Clear the audio input component so the user can speak again, and play the response audio
    audio_out_update = gr.update(value=audio_path_out, visible=True) if audio_path_out else gr.update()
    
    return new_history, None, audio_out_update


def respond(message, history):
    """Process a user message through the SHL agent pipeline."""
    if not message or not message.strip():
        return "Please type a message to get started."

    # Format current history into App Message models
    agent_messages = []

    # Convert Gradio format back to Message models
    for turn in history:
        role = turn["role"]
        if role in ["user", "assistant"]:
            agent_messages.append(Message(role=role, content=turn["content"]))

    # Append the new user message
    agent_messages.append(Message(role="user", content=message))

    try:
        agent = get_agent()
        response = agent.process(agent_messages)

        reply_content = response.reply

        # Format recommendations if present
        if response.recommendations:
            reply_content += "\n\n**Recommended Assessments:**\n"
            for idx, rec in enumerate(response.recommendations, 1):
                reply_content += f"{idx}. **[{rec.name}]({rec.url})** — Type: `{rec.test_type}`\n"

        if response.end_of_conversation:
            reply_content += "\n\n---\n✅ *Conversation complete. Thank you for using the SHL Assessment Recommender!*"

        return reply_content
    except Exception as e:
        print(f"[App] Error in respond: {e}")
        import traceback
        traceback.print_exc()
        return f"Sorry, something went wrong: {str(e)}"


def strip_markdown(text):
    """Remove markdown formatting for plain-text output."""
    text = re.sub(r'\*\*(.+?)\*\*', r'\1', text)  # bold
    text = re.sub(r'\*(.+?)\*', r'\1', text)        # italic
    text = re.sub(r'`(.+?)`', r'\1', text)          # inline code
    text = re.sub(r'\[(.+?)\]\((.+?)\)', r'\1 (\2)', text)  # links
    text = re.sub(r'^#{1,6}\s+', '', text, flags=re.MULTILINE)  # headings
    text = re.sub(r'^---+$', '', text, flags=re.MULTILINE)  # hr
    return text.strip()


def sanitize_for_pdf(text):
    """Make text safe for fpdf Helvetica (latin-1 only)."""
    # Replace common Unicode chars with ASCII equivalents
    replacements = {
        '\u2014': '-',   # em-dash
        '\u2013': '-',   # en-dash
        '\u2018': "'",   # left single quote
        '\u2019': "'",   # right single quote
        '\u201c': '"',   # left double quote
        '\u201d': '"',   # right double quote
        '\u2026': '...',  # ellipsis
        '\u2022': '*',   # bullet
        '\u00a0': ' ',   # non-breaking space
    }
    for orig, repl in replacements.items():
        text = text.replace(orig, repl)
    # Remove any remaining emojis/special chars
    text = text.encode('latin-1', 'replace').decode('latin-1')
    return text


def export_chat_to_pdf(history):
    """Export the chat history to a downloadable PDF file."""
    if not history or len(history) == 0:
        return None

    try:
        from fpdf import FPDF

        pdf = FPDF()
        pdf.set_auto_page_break(auto=True, margin=20)
        pdf.add_page()

        # Title
        pdf.set_font("Helvetica", "B", 18)
        pdf.cell(0, 12, "SHL Assessment Recommender", new_x="LMARGIN", new_y="NEXT", align="C")
        pdf.set_font("Helvetica", "", 10)
        pdf.set_text_color(120, 120, 120)
        pdf.cell(0, 8, sanitize_for_pdf(f"Chat Transcript - {datetime.now().strftime('%B %d, %Y at %I:%M %p')}"), new_x="LMARGIN", new_y="NEXT", align="C")
        pdf.ln(6)
        pdf.set_draw_color(200, 200, 200)
        pdf.line(10, pdf.get_y(), 200, pdf.get_y())
        pdf.ln(6)

        # Chat messages
        for turn in history:
            role = turn.get("role", "unknown")
            content = turn.get("content", "")
            clean_content = strip_markdown(content)

            if role == "user":
                pdf.set_font("Helvetica", "B", 11)
                pdf.set_text_color(30, 100, 200)
                pdf.cell(0, 7, "You:", new_x="LMARGIN", new_y="NEXT")
            else:
                pdf.set_font("Helvetica", "B", 11)
                pdf.set_text_color(34, 139, 34)
                pdf.cell(0, 7, "SHL Assistant:", new_x="LMARGIN", new_y="NEXT")

            pdf.set_font("Helvetica", "", 10)
            pdf.set_text_color(40, 40, 40)
            safe_text = sanitize_for_pdf(clean_content)
            pdf.multi_cell(0, 6, safe_text)
            pdf.ln(4)

        # Footer
        pdf.ln(6)
        pdf.set_draw_color(200, 200, 200)
        pdf.line(10, pdf.get_y(), 200, pdf.get_y())
        pdf.ln(4)
        pdf.set_font("Helvetica", "I", 8)
        pdf.set_text_color(150, 150, 150)
        pdf.cell(0, 6, sanitize_for_pdf("Generated by SHL Assessment Recommender - Powered by NVIDIA NIM"), new_x="LMARGIN", new_y="NEXT", align="C")

        # Save to temp file
        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf", prefix="shl_chat_")
        pdf.output(tmp.name)
        return tmp.name

    except Exception as e:
        print(f"[PDF] Export failed: {e}")
        import traceback
        traceback.print_exc()
        return None


# ─── Custom CSS ───
css = """
/* Constrain SVG/avatar icons */
.avatar-container, .avatar {
    max-width: 36px !important;
    max-height: 36px !important;
}
.avatar-container svg, .avatar-container img, .avatar svg, .avatar img {
    width: 100% !important;
    height: 100% !important;
    max-width: 36px !important;
    max-height: 36px !important;
    object-fit: contain !important;
}
/* Don't constrain images inside chat messages (e.g. recommendation cards) */
.message img, .bot img, .user img {
    max-width: 100% !important;
    max-height: none !important;
}
/* Make chatbot taller */
.chatbot {
    min-height: 500px !important;
}
/* Style the header row */
.header-row {
    display: flex;
    justify-content: space-between;
    align-items: center;
}
"""


# ─── Build the Gradio UI ───
head_js = """
<script>
window.voiceModeActive = false;
window.recognitionInstance = null;
window.lastPlayedSrc = "";

function toggleVoiceMode() {
    window.voiceModeActive = !window.voiceModeActive;
    console.log("[Voice] Voice mode active:", window.voiceModeActive);
    
    const statusBtn = document.querySelector("#voice_mode_btn");
    if (window.voiceModeActive) {
        if (statusBtn) {
            statusBtn.textContent = "🟢 Voice Mode: ACTIVE";
            statusBtn.style.backgroundColor = "#22C55E";
            statusBtn.style.color = "white";
        }
        startListening();
    } else {
        if (statusBtn) {
            statusBtn.textContent = "🎙️ Start Real-time Voice Chat";
            statusBtn.style.backgroundColor = "";
            statusBtn.style.color = "";
        }
        stopListening();
    }
}

function startListening() {
    if (!window.voiceModeActive) return;
    
    const SpeechRecognition = window.SpeechRecognition || window.webkitSpeechRecognition;
    if (!SpeechRecognition) {
        alert("Web Speech API is not supported in this browser. Please use Chrome, Safari, or Edge.");
        window.voiceModeActive = false;
        return;
    }
    
    if (window.recognitionInstance) {
        try { window.recognitionInstance.stop(); } catch(e) {}
    }
    
    const recognition = new SpeechRecognition();
    recognition.continuous = false;
    recognition.interimResults = false;
    recognition.lang = 'en-US';
    window.recognitionInstance = recognition;
    
    recognition.onstart = function() {
        console.log("[Voice] Listening...");
        const statusBtn = document.querySelector("#voice_mode_btn");
        if (statusBtn) statusBtn.textContent = "🔴 Listening to your voice...";
    };
    
    recognition.onresult = function(event) {
        const transcript = event.results[0][0].transcript;
        console.log("[Voice] Transcribed:", transcript);
        
        const textbox = document.querySelector("#user_msg_input textarea");
        const sendBtn = document.querySelector("#send_btn");
        
        if (textbox && sendBtn && transcript.trim()) {
            textbox.value = transcript;
            textbox.dispatchEvent(new Event('input', { bubbles: true }));
            setTimeout(() => {
                sendBtn.click();
            }, 100);
        }
    };
    
    recognition.onerror = function(event) {
        console.error("[Voice] Error:", event.error);
        if (event.error === 'no-speech') {
            if (window.voiceModeActive) {
                setTimeout(startListening, 300);
            }
        }
    };
    
    recognition.onend = function() {
        console.log("[Voice] Stopped listening.");
        const statusBtn = document.querySelector("#voice_mode_btn");
        if (statusBtn && window.voiceModeActive) {
            statusBtn.textContent = "🟢 Voice Mode: ACTIVE";
        }
    };
    
    recognition.start();
}

function stopListening() {
    if (window.recognitionInstance) {
        try { window.recognitionInstance.stop(); } catch(e) {}
    }
}

// Global capture for Audio playing events to toggle microphone
document.addEventListener("play", (event) => {
    if (event.target.tagName === "AUDIO") {
        console.log("[Voice] Audio started playing. Stopping microphone.");
        stopListening();
    }
}, true);

document.addEventListener("ended", (event) => {
    if (event.target.tagName === "AUDIO") {
        console.log("[Voice] Audio finished playing. Resuming microphone.");
        if (window.voiceModeActive) {
            setTimeout(startListening, 400);
        }
    }
}, true);

// User interruption via click anywhere on the chatbot or clicking the voice button
document.addEventListener("click", (event) => {
    const audioEl = document.querySelector("#audio_output audio");
    if (audioEl && !audioEl.paused) {
        const isChat = event.target.closest(".chatbot") || event.target.closest("#voice_mode_btn") || event.target.closest("#user_msg_input");
        if (isChat) {
            console.log("[Voice] User interrupted assistant via click. Stopping audio and starting listening.");
            audioEl.pause();
            audioEl.currentTime = 0;
            if (window.voiceModeActive) {
                startListening();
            }
        }
    }
});

// User interruption via Spacebar keypress
document.addEventListener("keydown", (event) => {
    if (event.code === "Space") {
        const audioEl = document.querySelector("#audio_output audio");
        if (audioEl && !audioEl.paused) {
            console.log("[Voice] User interrupted assistant via Space key. Stopping audio and starting listening.");
            audioEl.pause();
            audioEl.currentTime = 0;
            event.preventDefault();
            if (window.voiceModeActive) {
                startListening();
            }
        }
    }
});

// MutationObserver to guarantee autoplay immediately when Gradio updates the audio source
function setupAudioAutoplayObserver() {
    const interval = setInterval(() => {
        const target = document.querySelector("#audio_output");
        if (target) {
            clearInterval(interval);
            console.log("[Voice] Found audio output container. Setting up MutationObserver.");
            const observer = new MutationObserver((mutations) => {
                mutations.forEach((mutation) => {
                    const audioEl = document.querySelector("#audio_output audio");
                    if (audioEl && audioEl.src && audioEl.src !== window.lastPlayedSrc) {
                        console.log("[Voice] Auto-playing new audio source:", audioEl.src);
                        window.lastPlayedSrc = audioEl.src;
                        audioEl.play().catch(e => console.error("[Voice] Play failed:", e));
                    }
                });
            });
            observer.observe(target, { childList: true, subtree: true, attributes: true });
        }
    }, 1000);
}

if (document.readyState === "loading") {
    document.addEventListener("DOMContentLoaded", setupAudioAutoplayObserver);
} else {
    setupAudioAutoplayObserver();
}
</script>
"""

with gr.Blocks(
    title="SHL Assessment Recommender",
    css=css,
    head=head_js,
    theme=gr.themes.Soft(
        primary_hue="green",
        neutral_hue="gray",
    )
) as demo:

    # Header with action buttons
    with gr.Row():
        with gr.Column(scale=7):
            gr.Markdown("# 🧠 SHL Assessment Recommender")
            gr.Markdown("A conversational AI to help you find the right SHL assessments. Just chat naturally or use Real-time Voice!")
        with gr.Column(scale=3, min_width=280):
            with gr.Row():
                tts_btn = gr.Button("🔊 Listen", variant="secondary", size="sm")
                export_btn = gr.Button("📄 Export PDF", variant="secondary", size="sm")

    # Quick-start guide
    with gr.Accordion("📖 How to use this tool", open=False):
        gr.Markdown("""
**Step 1:** Tell the assistant about the role you're hiring for (e.g. *"I'm hiring a Java developer"*).

**Step 2:** Answer follow-up questions — the assistant will ask about experience level, skills, and preferences.

**Step 3:** Get personalized SHL assessment recommendations with direct links.

**Step 4:** Refine your list — ask to add, remove, or compare assessments.

**Step 5:** Export your conversation to PDF, listen to the response, or turn on **🎙️ Real-time Voice Chat** to talk hands-free!

> 💡 **Tip:** In Real-time Voice Chat mode, you just speak. The app will automatically transcribe your voice and answer back out loud, just like ChatGPT Voice Mode!
""")

    # Real-time Voice Control button
    with gr.Row():
        voice_mode_btn = gr.Button("🎙️ Start Real-time Voice Chat (Hands-free)", variant="primary", elem_id="voice_mode_btn")

    # Chat area
    chatbot = gr.Chatbot(
        value=[{
            "role": "assistant",
            "content": (
                "Hello! 👋 I'm your SHL Assessment advisor. "
                "Tell me about the role you're hiring for, and I'll recommend "
                "the best assessments from SHL's catalog.\n\n"
                "For example, you can say:\n"
                "- *\"I'm hiring a Java developer with 3 years experience\"*\n"
                "- *\"I need assessments for a senior project manager\"*\n"
                "- *\"Looking for cognitive ability tests for entry-level candidates\"*"
            )
        }],
        type="messages",
        label="SHL Chat Assistant",
        height=520,
        show_copy_button=True,
    )

    # Input area
    with gr.Row():
        with gr.Column(scale=8):
            msg = gr.Textbox(
                placeholder="Type your message here... (e.g. 'I want to hire a backend developer')",
                label="Your Message",
                show_label=False,
                scale=9,
                lines=1,
                max_lines=3,
                elem_id="user_msg_input",
            )
        with gr.Column(scale=2, min_width=100):
            submit_btn = gr.Button("Send", variant="primary", scale=1, elem_id="send_btn")

    # Voice input component
    with gr.Row():
        audio_in = gr.Audio(
            sources=["microphone"],
            type="filepath",
            label="🎤 Or record / upload audio manually",
            show_download_button=False,
        )

    # Clear button
    clear = gr.Button("🗑️ Clear Chat History", variant="stop", size="sm")

    # Hidden outputs
    pdf_output = gr.File(label="Download PDF", visible=False)
    audio_output = gr.Audio(label="🔊 Assistant Voice", visible=False, autoplay=True, elem_id="audio_output")

    # ─── Event Handlers ───
    def user_msg(user_message, history):
        if not user_message or not user_message.strip():
            return "", history
        return "", history + [{"role": "user", "content": user_message.strip()}]

    def bot_msg(history):
        if not history or history[-1]["role"] != "user":
            return history, gr.update()
        user_message = history[-1]["content"]
        history_before = history[:-1]

        bot_response = respond(user_message, history_before)
        history.append({"role": "assistant", "content": bot_response})
        
        # Auto TTS generation for seamless voice replies
        audio_path = text_to_speech(history)
        audio_update = gr.update(value=audio_path, visible=True) if audio_path else gr.update()
        return history, audio_update

    def handle_export(history):
        pdf_path = export_chat_to_pdf(history)
        if pdf_path:
            return gr.update(value=pdf_path, visible=True)
        return gr.update(visible=False)

    # Send on Enter
    msg.submit(user_msg, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot_msg, chatbot, [chatbot, audio_output]
    )

    # Send on button click
    submit_btn.click(user_msg, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot_msg, chatbot, [chatbot, audio_output]
    )

    # Voice STT Transcription and Response event handler
    audio_in.change(
        transcribe_and_respond,
        [audio_in, chatbot],
        [chatbot, audio_in, audio_output],
        queue=False
    )

    # Voice Mode Toggle handler (JS-only)
    voice_mode_btn.click(None, None, None, js="() => { toggleVoiceMode(); }")

    # Export PDF
    export_btn.click(handle_export, chatbot, pdf_output)

    # Text-to-Speech
    def handle_tts(history):
        audio_path = text_to_speech(history)
        if audio_path:
            return gr.update(value=audio_path, visible=True)
        return gr.update(visible=False)

    tts_btn.click(handle_tts, chatbot, audio_output)

    # Clear chat
    def reset_chat():
        return [{
            "role": "assistant",
            "content": (
                "Hello! 👋 I'm your SHL Assessment advisor. "
                "Tell me about the role you're hiring for, and I'll recommend "
                "the best assessments from SHL's catalog.\n\n"
                "For example, you can say:\n"
                "- *\"I'm hiring a Java developer with 3 years experience\"*\n"
                "- *\"I need assessments for a senior project manager\"*\n"
                "- *\"Looking for cognitive ability tests for entry-level candidates\"*"
            )
        }], gr.update(visible=False), gr.update(visible=False), None

    clear.click(reset_chat, None, [chatbot, pdf_output, audio_output, audio_in], queue=False)


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)