import gradio as gr
import os
from groq import Groq

def initialize_groq_client():
    api_key = os.getenv("GROQ_API_KEY")
    if not api_key:
        raise ValueError("GROQ_API_KEY environment variable not set")
    return Groq(api_key=api_key)

def translate_to_swahili(text, temperature=0.3):
    if not text.strip():
        return "Please enter some text to translate."
    
    try:
        client = initialize_groq_client()
        prompt = f"""You are a professional translator. Translate the following English text to Swahili. 
Provide only the Swahili translation, nothing else.
English text: {text}
Swahili translation:"""

        completion = client.chat.completions.create(
            model="meta-llama/llama-4-scout-17b-16e-instruct",
            messages=[
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            temperature=temperature,
            max_completion_tokens=1024,
            top_p=1,
            stream=False,
            stop=None,
        )
        
        return completion.choices[0].message.content.strip()
        
    except Exception as e:
        return f"Error during translation: {str(e)}"

def translate_with_streaming(text, temperature=0.3):
    if not text.strip():
        yield "Please enter some text to translate."
        return
    
    try:
        client = initialize_groq_client()
        prompt = f"""You are a professional translator. Translate the following English text to Swahili. 
Provide only the Swahili translation, nothing else.
English text: {text}
Swahili translation:"""

        completion = client.chat.completions.create(
            model="meta-llama/llama-4-scout-17b-16e-instruct",
            messages=[
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            temperature=temperature,
            max_completion_tokens=1024,
            top_p=1,
            stream=True,
            stop=None,
        )
        
        full_response = ""
        for chunk in completion:
            if chunk.choices[0].delta.content:
                full_response += chunk.choices[0].delta.content
                yield full_response
                
    except Exception as e:
        yield f"Error during translation: {str(e)}"

def create_interface():
    with gr.Blocks(title="English to Swahili Translator", theme=gr.themes.Soft()) as iface:
        gr.Markdown("# 🌍 English to Swahili Translator")
        gr.Markdown("Powered by Groq's Llama 4 - Translate English text to Swahili")
        
        with gr.Row():
            with gr.Column(scale=1):
                input_text = gr.Textbox(
                    label="English Text",
                    placeholder="Enter English text to translate...",
                    lines=5,
                    max_lines=10
                )
                
                temperature = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.3,
                    step=0.1,
                    label="Temperature (creativity level)",
                    info="Lower values = more consistent, Higher values = more creative"
                )
                
                with gr.Row():
                    translate_btn = gr.Button("🔄 Translate", variant="primary")
                    clear_btn = gr.Button("🗑️ Clear")
            
            with gr.Column(scale=1):
                output_text = gr.Textbox(
                    label="Swahili Translation",
                    lines=5,
                    max_lines=10,
                    interactive=False
                )
        
        gr.Examples(
            examples=[
                ["Hello, how are you?"],
                ["I love learning new languages."],
                ["The weather is beautiful today."],
                ["Thank you for your help."],
                ["What time is it?"]
            ],
            inputs=[input_text],
            label="Try these examples:"
        )
        
        translate_btn.click(
            fn=translate_with_streaming,
            inputs=[input_text, temperature],
            outputs=[output_text],
            show_progress=True
        )
        
        clear_btn.click(
            fn=lambda: ("", ""),
            outputs=[input_text, output_text]
        )
        
        input_text.submit(
            fn=translate_with_streaming,
            inputs=[input_text, temperature],
            outputs=[output_text],
            show_progress=True
        )
    
    return iface

if __name__ == "__main__":
    iface = create_interface()
    iface.launch(
        debug=True,
        share=False,
        server_name="0.0.0.0",
        server_port=7860
    )

def launch_hf_space():
    iface = create_interface()
    iface.launch()