import gradio as gr
import os
import requests
from gtts import gTTS
import tempfile
import uuid
import datetime
from fpdf import FPDF
import whisper

# Whisper for transcription
model = whisper.load_model("base")

# OpenRouter chat
def ai_chat(text):
    url = "https://openrouter.ai/api/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {os.getenv('OPENROUTER_API_KEY')}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": "mistralai/mistral-7b-instruct",
        "messages": [{"role": "user", "content": text}]
    }
    res = requests.post(url, headers=headers, json=payload)
    if res.ok:
        return res.json()["choices"][0]["message"]["content"]
    return "[OpenRouter error]"

# PDF generation
def generate_pdf(chat_history, session_id):
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)
    for turn in chat_history:
        pdf.multi_cell(0, 10, f"You: {turn['user']}\nAI: {turn['ai']}\n")
    path = f"{session_id}_chat.pdf"
    pdf.output(path)
    return path

# Main voice chat function
def voice_chat(audio_file, chat_history, session_id):
    if not audio_file:
        return None, chat_history, session_id, None

    if not session_id:
        session_id = str(uuid.uuid4())

    result = model.transcribe(audio_file)
    user_text = result["text"]
    ai_reply = ai_chat(user_text)

    # TTS
    tts = gTTS(ai_reply, lang='en')
    tts_path = tempfile.NamedTemporaryFile(suffix=".mp3", delete=False)
    tts.save(tts_path.name)

    # Update chat history
    chat_history.append({"user": user_text, "ai": ai_reply})

    # Build display text
    display = ""
    for turn in chat_history:
        display += f"🧑‍💬 **You**: {turn['user']}\n\n🤖 **AI**: {turn['ai']}\n\n"

    return tts_path.name, chat_history, session_id, display

# PDF download
def download_pdf(chat_history, session_id):
    if not chat_history or not session_id:
        return None
    return generate_pdf(chat_history, session_id)

# === UI ===
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 🎤 Chat with AI (Voice + Text, Powered by OpenRouter)")

    session_id = gr.State()
    chat_history = gr.State([])

    chatbox = gr.Markdown("")

    with gr.Row():
        mic = gr.Audio(label="", type="filepath", show_label=False, interactive=True)
        send_btn = gr.Button("🎙️ Speak", variant="primary")

    with gr.Row():
        audio_out = gr.Audio(label="🔊 AI Voice Reply", interactive=False)

    with gr.Row():
        download_btn = gr.Button("📥 Download Chat as PDF")
        download_link = gr.File()

    send_btn.click(voice_chat, inputs=[mic, chat_history, session_id],
                   outputs=[audio_out, chat_history, session_id, chatbox])

    download_btn.click(download_pdf, inputs=[chat_history, session_id], outputs=[download_link])

demo.launch()