import os
import gradio as gr
from huggingface_hub import InferenceClient

# 1. Securely grab your token from Space Secrets 🔒
HF_TOKEN = os.getenv("HF_TOKEN")

# 2. Initialize the Client with Llama-3 🧠
client = InferenceClient(model="meta-llama/Meta-Llama-3-8B-Instruct", token=HF_TOKEN)

# --- AGENT 1: STATELESS LOGIC ---
def respond_stateless(message):
    # Only the current message is sent to the AI
    messages = [
        {"role": "system", "content": "You are a stateless AI. Answer only the current prompt."},
        {"role": "user", "content": message}
    ]
    
    response = ""
    for message_obj in client.chat_completion(messages, max_tokens=512, stream=True):
        if hasattr(message_obj, 'choices') and len(message_obj.choices) > 0:
            content = message_obj.choices[0].delta.content
            if content:
                response += content
    return response

# --- AGENT 2: STATEFUL LOGIC ---
def respond_stateful(message, history):
    # 'history' is our internal "Hidden Notebook" (list of tuples)
    messages = [{"role": "system", "content": "You are a Context-Aware Assistant. Use the history to remember details."}]
    
    # Rebuild the full context for the AI 📑
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})
    
    messages.append({"role": "user", "content": message})
    
    response = ""
    for message_obj in client.chat_completion(messages, max_tokens=512, stream=True):
        if hasattr(message_obj, 'choices') and len(message_obj.choices) > 0:
            content = message_obj.choices[0].delta.content
            if content:
                response += content
            
    # Update our internal tuple-based state
    history.append((message, response))
    return response, history

# --- COORDINATOR FUNCTION ---
def chat_coordinator(user_message, chat1_ui_data, state_history):
    # Process Stateless Agent
    res1 = respond_stateless(user_message)
    # Convert to dictionary format for UI 1
    chat1_ui_data.append({"role": "user", "content": user_message})
    chat1_ui_data.append({"role": "assistant", "content": res1})
    
    # Process Stateful Agent
    res2, updated_state = respond_stateful(user_message, state_history)
    
    # Convert our internal "Hidden Notebook" (tuples) into UI dictionaries for UI 2 🔄
    chat2_ui_list = []
    for u, a in updated_state:
        chat2_ui_list.append({"role": "user", "content": u})
        chat2_ui_list.append({"role": "assistant", "content": a})
    
    # Returns: (Clear input, Agent 1 UI, Agent 2 UI, Internal State, Counter)
    return "", chat1_ui_data, chat2_ui_list, updated_state, len(updated_state)

# --- GRADIO UI LAYOUT ---
with gr.Blocks() as demo:
    gr.Markdown("Tabula Rasa vs Total Recall")
    
    # Internal state for the 'Hidden Notebook' logic
    hidden_notebook = gr.State([])

    # Main horizontal split: Sidebar (1) and Workspace (4)
    with gr.Row():
        
        # --- LEFT SIDEBAR (System Controls) ---
        with gr.Column(scale=1, variant="panel"):
            gr.Markdown("### ⚙️ System Controls")
            mem_counter = gr.Number(value=0, label="Memory Context (Turns)")
            
            # A 'stop' variant makes the reset button red/distinct
            reset_btn = gr.Button("🗑️ Reset Stateful Memory", variant="stop")
            
            gr.Markdown("---")
            gr.Info("Agent 2 rebuilds its context from the Hidden Notebook for every message.")

        # --- RIGHT MAIN PANEL (Agent Comparison) ---
        with gr.Column(scale=4):
            with gr.Row():
                # Column for Agent 1
                with gr.Column():
                    gr.Markdown("### Agent Tabula Rasa")
                    gr.Markdown("_Stateless: This agent has no memory but knows everything")
                    chat1_ui = gr.Chatbot(height=450, show_label=False)
                
                # Column for Agent 2
                with gr.Column():
                     gr.Markdown("### Agent Total Recall")
                     gr.Markdown("_Stateful: This agent knows and remembers everything")
                     chat2_ui = gr.Chatbot(height=450, show_label=False)
            
            # User Input Bar at the bottom
            with gr.Row():
                user_input = gr.Textbox(
                    show_label=False, 
                    placeholder="Ask both agents a question to compare memory...", 
                    scale=4
                )
                submit_btn = gr.Button("Send 🚀", variant="primary", scale=1)

    # --- WIRING IT ALL TOGETHER ---
    config = {
        "fn": chat_coordinator,
        "inputs": [user_input, chat1_ui, hidden_notebook],
        "outputs": [user_input, chat1_ui, chat2_ui, hidden_notebook, mem_counter]
    }
    
    user_input.submit(**config)
    submit_btn.click(**config)
    
    # Resetting clears the UI, the state, and the turn counter
    reset_btn.click(
        lambda: ([], [], [], 0), 
        None, 
        [chat1_ui, chat2_ui, hidden_notebook, mem_counter]
    )

# Final launch with a soft theme
demo.launch(theme=gr.themes.Soft())