import gradio as gr
from huggingface_hub import InferenceClient

# ⚠️ PASTE YOUR REAL TOKEN HERE
HF_TOKEN = "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..."  

# ✅ USE A DIFFERENT MODEL THAT WORKS BETTER ON SPACES
client = InferenceClient(
    model="google/gemma-2-2b-it",  
    token=HF_TOKEN
)

def chat_with_ai(prompt, history):
    if not prompt.strip():
        return "", history

    try:
        # Using a simpler prompt format for Gemma
        full_prompt = f"<start_of_turn>user\nYou are Hermus, an AI assistant. {prompt}<end_of_turn>\n<start_of_turn>model\n"
        
        response = client.text_generation(
            full_prompt,
            max_new_tokens=512,
            temperature=0.7
        )
        reply = response.strip()
        
    except Exception as e:
        reply = f"⚠️ Error: {str(e)}. If this fails, we will switch to a completely offline mode next."

    history.append({"role": "user", "content": prompt})
    history.append({"role": "assistant", "content": reply})
    return "", history

# UI
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 Hermus AI Project Assistant (Gemma Version)")
    
    chatbot = gr.Chatbot(height=400)
    msg = gr.Textbox(label="Type your request here")
    clear = gr.Button("Clear Chat")
    
    msg.submit(chat_with_ai, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

demo.launch(theme="soft")