import gradio as gr
import whisper
import os
from groq import Groq
import tempfile
from gtts import gTTS
import fitz  # PyMuPDF for PDF reading

# Load Whisper model
whisper_model = whisper.load_model("base")

# Groq API setup
groq_api_key = "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"
client = Groq(api_key=groq_api_key)

# Urdu legal Q&A (text input)
def urdu_chat(question):
    response = client.chat.completions.create(
        model="llama3-70b-8192",
        messages=[
            {"role": "system", "content": "You are a helpful AI assistant that answers legal questions in Urdu simply."},
            {"role": "user", "content": question}
        ]
    )
    return response.choices[0].message.content

# Urdu legal Q&A (voice input)
def urdu_chat_voice(audio_file_path):
    result = whisper_model.transcribe(audio_file_path, language="ur")
    question = result["text"]

    response = client.chat.completions.create(
        model="llama3-70b-8192",
        messages=[
            {"role": "system", "content": "You are a helpful AI assistant that answers legal questions in Urdu simply."},
            {"role": "user", "content": question}
        ]
    )
    answer = response.choices[0].message.content

    tts = gTTS(answer, lang='ur')
    audio_path = tempfile.mktemp(suffix=".mp3")
    tts.save(audio_path)

    return question, answer, audio_path

# Urdu legal Q&A (PDF upload)
def analyze_pdf(pdf_file):
    try:
        doc = fitz.open(pdf_file)
        full_text = ""
        for page in doc:
            full_text += page.get_text()

        # Ask Groq for legal opinion
        response = client.chat.completions.create(
            model="llama3-70b-8192",
            messages=[
                {"role": "system", "content": "You are a legal expert AI. Analyze the following legal document in Urdu and give your expert opinion simply."},
                {"role": "user", "content": full_text}
            ]
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"⚠️ خرابی: {str(e)}"

# Gradio UI
with gr.Blocks(title="⚖ Urdu Legal Assistant") as demo:
    gr.Markdown("## 🧑‍⚖ اردو میں قانونی سوالات پوچھیں (ٹیکسٹ، آڈیو یا PDF سے)")

    with gr.Tab("✍ ٹیکسٹ چیٹ"):
        with gr.Row():
            text_input = gr.Textbox(lines=2, placeholder="اپنا قانونی سوال یہاں لکھیں...")
            text_output = gr.Textbox(label="جواب")
        text_btn = gr.Button("جواب حاصل کریں")
        text_btn.click(fn=urdu_chat, inputs=text_input, outputs=text_output)

    with gr.Tab("🎙 آواز سے پوچھیں"):
        with gr.Row():
            audio_input = gr.Audio(type="filepath", label="🔊 اپنی آواز میں سوال ریکارڈ کریں یا اپلوڈ کریں")
        with gr.Row():
            transcribed = gr.Textbox(label="📝 سوال (آڈیو سے متن)")
            ai_answer = gr.Textbox(label="🤖 AI جواب")
        with gr.Row():
            audio_response = gr.Audio(type="filepath", label="🔊 آڈیو میں جواب سنیں")
        voice_btn = gr.Button("جواب سنیں")
        voice_btn.click(fn=urdu_chat_voice, inputs=audio_input, outputs=[transcribed, ai_answer, audio_response])

    with gr.Tab("📄 قانونی PDF تجزیہ"):
        with gr.Row():
            pdf_input = gr.File(file_types=[".pdf"], label="اپنا قانونی PDF اپلوڈ کریں")
        with gr.Row():
            pdf_output = gr.Textbox(label="⚖ AI قانونی رائے", lines=10)
        pdf_btn = gr.Button("تجزیہ کریں")
        pdf_btn.click(fn=analyze_pdf, inputs=pdf_input, outputs=pdf_output)

demo.launch()