import os
import gradio as gr
from huggingface_hub import InferenceClient

# Authenticate with Hugging Face
client = InferenceClient(
    provider="sambanova",
    api_key=os.environ["HF_TOKEN"],
)

def llama4_chat(image_url, text_prompt):
    if not text_prompt:
        return "Please enter a text question or prompt."

    message_content = [{"type": "text", "text": text_prompt}]

    if image_url:
        message_content.append({
            "type": "image_url",
            "image_url": {"url": image_url}
        })

    messages = [{"role": "user", "content": message_content}]

    try:
        completion = client.chat.completions.create(
            model="meta-llama/Llama-4-Maverick-17B-128E-Instruct",
            messages=messages
        )
        return completion.choices[0].message.content
    except Exception as e:
        return f"Error: {e}"

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🦙 LLaMA 4 Chat (Text + Optional Image URL)")
    gr.Markdown("Enter a text prompt. Optionally provide an image URL.")

    with gr.Row():
        text_prompt_input = gr.Textbox(label="Text Prompt / Question", placeholder="What is happening in the image?", lines=2)
        image_url_input = gr.Textbox(label="Optional Image URL", placeholder="https://example.com/image.jpg")

    submit_btn = gr.Button("Generate Response")
    output_box = gr.Textbox(label="LLaMA 4 Response", lines=8)

    submit_btn.click(
        fn=llama4_chat,
        inputs=[image_url_input, text_prompt_input],
        outputs=output_box
    )

demo.launch()