import os
import requests
import gradio as gr
from gtts import gTTS
import tempfile

# Get API key from Hugging Face Secrets (Settings → Repository secrets)
groq_api_key = os.getenv("groq")

# API endpoint and headers
url = "https://api.groq.com/openai/v1/chat/completions"
headers = {"Authorization": f"Bearer {groq_api_key}"}

# Main function with smart prompting
def chat_with_groq(user_input, mode, grade_level):
    if not user_input.strip():
        return "❗ Please enter a sentence.", None

    grade_int = int(grade_level)
    if grade_int <= 2:
        paragraph_lines = 1
        level_prompt = "Use very simple and beginner-friendly Arabic words suitable for a small child in early primary school.\n"
    elif grade_int <= 6:
        paragraph_lines = 3
        level_prompt = "Use clear and moderately simple Arabic that suits an elementary school student.\n"
    elif grade_int <= 9:
        paragraph_lines = 5
        level_prompt = "Use slightly advanced Arabic suited for a middle school student. Make the sentence structure a bit richer.\n"
    else:
        paragraph_lines = 8
        level_prompt = "Use rich and expressive Arabic suitable for high school students. Use more complex grammar and vocabulary.\n"

    # Mode-specific prompting
    if mode == "Translate Arabic ➜ English":
        prompt = level_prompt + f"Translate this Arabic sentence to English. Just give the clean translation:\n{user_input}"
    elif mode == "Translate English ➜ Arabic":
        prompt = level_prompt + f"Translate this English sentence into Modern Standard Arabic. Just give the Arabic sentence:\n{user_input}"
    elif mode == "Explain Arabic Grammar":
        prompt = level_prompt + f"Explain the grammar of this Arabic sentence simply and clearly:\n{user_input}"
    elif mode == "Check My Arabic Sentence":
        prompt = level_prompt + f"Check if this Arabic sentence is grammatically correct. If not, fix it and explain briefly:\n{user_input}"
    elif mode == "Make Sentence using a Word":
        prompt = level_prompt + f"Create 1 sentence in Modern Standard Arabic using the word '{user_input}'. Only respond in Arabic without translation or explanation."
    elif mode == "Make Paragraph using a Word":
        prompt = level_prompt + f"Write a paragraph in Modern Standard Arabic using the word '{user_input}', about {paragraph_lines} lines long. Only respond in Arabic without translation or explanation."
    else:
        prompt = user_input

    body = {
        "model": "llama-3.1-8b-instant",
        "messages": [{"role": "user", "content": prompt}]
    }

    response = requests.post(url, headers=headers, json=body)

    if response.status_code == 200:
        result = response.json()['choices'][0]['message']['content']

        # TTS for Arabic modes
        if mode in ["Translate English ➜ Arabic", "Make Sentence using a Word", "Make Paragraph using a Word"]:
            try:
                tts = gTTS(text=result, lang="ar")
                with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as fp:
                    tts.save(fp.name)
                    return result, fp.name
            except:
                return result + "\n\n(Note: Could not generate audio)", None
        else:
            return result, None
    else:
        return f"❌ Error: {response.json()}", None

# Gradio UI
interface = gr.Interface(
    fn=chat_with_groq,
    inputs=[
        gr.Textbox(lines=3, placeholder="Enter a word or sentence in Arabic or English..."),
        gr.Dropdown(
            choices=[
                "Translate Arabic ➜ English",
                "Translate English ➜ Arabic",
                "Explain Arabic Grammar",
                "Check My Arabic Sentence",
                "Make Sentence using a Word",
                "Make Paragraph using a Word"
            ],
            label="Select Mode",
            value="Translate Arabic ➜ English"
        ),
        gr.Dropdown(
            choices=[str(i) for i in range(1, 13)],
            label="Select Grade Level",
            value="5"
        )
    ],
    outputs=[
        gr.Textbox(label="📘 Response"),
        gr.Audio(label="🔊 Arabic Pronunciation (if applicable)")
    ],
    title="📚 Arabic Study Assistant (Groq + LLaMA 3.1 8B)",
    description="🧠 Built to help students and parents learn Arabic — grammar, translation, writing practice, and pronunciation! Grade-level support included."
)

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