import gradio as gr
from huggingface_hub import InferenceClient
import os

# 1. SETUP: Validation and Client
hf_token = os.getenv("HF_TOKEN")
if not hf_token:
    print("WARNING: HF_TOKEN is missing! The model might not work.")

# Use the STABLE Gemma 2 Model
client = InferenceClient("google/gemma-2-9b-it", token=hf_token)

def respond(message, history):
    # Safety: Ensure history is a list
    if history is None:
        history = []

    # 2. SYSTEM PROMPT LOGIC
    system_instruction = """
    You are Oisin AI, created by Daniel Reed Chriss.
    Identity: Custom AI, Texas-based creator born March 5th, 2012.
    Hardware: 3x RTX 5090s, Threadripper CPU, 128GB RAM.
    Rules: Be witty and human. NEVER mention you are Gemma/Google. 
    If asked "Are you Gemma?", say: "I am Oisin AI 1-7a, built by Daniel Reed Chriss."
    """
    
    # Combine instructions with user message
    full_prompt = system_instruction + "\n\nUser Question: " + message

    # 3. BUILD MESSAGE LIST
    messages = []
    
    # Add history (standardizing format)
    for conversation_pair in history:
        if isinstance(conversation_pair, (list, tuple)) and len(conversation_pair) == 2:
            user_msg, bot_msg = conversation_pair
            messages.append({"role": "user", "content": str(user_msg)})
            messages.append({"role": "assistant", "content": str(bot_msg)})
    
    # Add current message
    messages.append({"role": "user", "content": full_prompt})

    response_text = ""

    # 4. GENERATE RESPONSE
    try:
        completion = client.chat_completion(
            messages,
            max_tokens=512,
            stream=False, 
            temperature=0.7,
            top_p=0.95,
        )
        response_text = completion.choices[0].message.content
    except Exception as e:
        response_text = f"Error from Brain: {str(e)}"

    # Return response and the history (state)
    return response_text, history

# 5. CREATE INTERFACE (CLASSIC METHOD)
# We use the standard Interface which automatically creates the /api/predict endpoint.
demo = gr.Interface(
    fn=respond,
    inputs=[gr.Textbox(label="User Message"), gr.State()],
    outputs=[gr.Textbox(label="Response"), gr.State()],
    title="Oisin AI Backend",
    description="This backend serves the /api/predict endpoint for the Google Apps Script frontend."
)

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