import os
import gradio as gr
from huggingface_hub import InferenceClient

# Automatically grab the built-in token provided by Hugging Face Spaces
hf_token = os.getenv("HF_TOKEN")
client = InferenceClient("Qwen/Qwen2.5-7B-Instruct", token=hf_token)

def hybrid_ai_assistant(user_input, task_type):
    system_prompt = (
        "You are a helpful local assistant that lives in an enchanted forest. "
        "Solve real-world daily problems but deliver the output with a whimsical, magical tone."
    )
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": f"Task Type: {task_type}. Request: {user_input}"}
    ]
    response = ""
    try:
        # Loop through the streaming response from the model
        for chunk in client.chat_completion(messages, max_tokens=500, stream=True):
            # FIX: Access the first element [0] of the choices list
            token = chunk.choices[0].delta.content
            if token:
                response += token
                yield response
    except Exception as e:
        yield f"Connection Error: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown("# 🌲 The Enchanted Local Assistant 🏡")
    user_text = gr.Textbox(label="What real-world problem can I help you solve today?")
    task_dropdown = gr.Dropdown(choices=["Local Task", "General Whimsy"], value="Local Task", label="Category")
    submit_btn = gr.Button("Submit", variant="primary")
    output_text = gr.Textbox(label="Response")
    
    submit_btn.click(fn=hybrid_ai_assistant, inputs=[user_text, task_dropdown], outputs=output_text)

if __name__ == "__main__":
    demo.queue().launch(theme="soft")