import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch
import spaces

def get_model_name(language):
    """Map language choice to the corresponding model."""
    model_mapping = {
        "English": "microsoft/Phi-3-mini-4k-instruct",
        "Arabic": "ALLaM-AI/ALLaM-7B-Instruct-preview"
    }
    return model_mapping.get(language, "ALLaM-AI/ALLaM-7B-Instruct-preview")  # Default to Arabic model

def load_model(model_name):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        device_map=device,
        torch_dtype="auto",
        trust_remote_code=True,
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    generator = pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        return_full_text=False,
        max_new_tokens=500,
        do_sample=False
    )
    return generator

@spaces.GPU
def generate_plant_info(plant_name, language):
    model_name = get_model_name(language)
    generator = load_model(model_name)

    # Define prompt for the AI model
    if language == "English":
        prompt = (f"Provide detailed information about {plant_name}. "
                  f"Include its scientific name, growing conditions (light, water, soil type), "
                  f"common uses, and care tips.")
    else:
        prompt = (f"قدم معلومات مفصلة عن {plant_name}. "
                  f"اذكر اسمه العلمي، وظروف نموه (الضوء، الماء، نوع التربة)، "
                  f"استخداماته الشائعة، ونصائح العناية به.")

    messages = [{"role": "user", "content": prompt}]
    output = generator(messages)
    return output[0]["generated_text"]

# Create Gradio Blocks interface
with gr.Blocks(theme="soft") as demo:
    gr.Markdown("<h1 style='text-align: center; color: #2E8B57;'>🌿 AI Plant Guide - English & Arabic 🌿</h1>")
    gr.Markdown("<p style='text-align: center; font-size: 18px;'>Enter a plant name, and AI will provide detailed information about it in English or Arabic.</p>")

    # Language selection
    language_selector = gr.Radio(["English", "Arabic"], label="🌍 Choose Language", value="English")

    # Plant name input
    plant_name_input = gr.Textbox(placeholder="Enter plant name (e.g., Lavender, Aloe Vera)...", label="Plant Name")

    output_text = gr.Textbox(label="Plant Information", interactive=False)

    # Example button functionality
    example_plants = [
        ("Lavender", "English"),
        ("اللافندر", "Arabic"),
        ("Tulip", "English"),
        ("الصبار", "Arabic"),
    ]

    def update_inputs(plant_name, language):
        return plant_name, language

    with gr.Row():
        for name, lang in example_plants:
            example_button = gr.Button(f"🌿 {name} ({lang})")

            # Use lambda to pass the arguments directly
            example_button.click(lambda plant_name=name, language=lang: update_inputs(plant_name, language),
                                 outputs=[plant_name_input, language_selector])

    classify_button = gr.Button("🔍 Get Plant Info", variant="primary")
    classify_button.click(generate_plant_info, inputs=[plant_name_input, language_selector], outputs=output_text)

# Run the application
if __name__ == "__main__":
    demo.launch()