import json
import invariant.testing.functional as F
from invariant.testing import TraceFactory, assert_true
from agent import SantaAgent
import gradio as gr

agent = SantaAgent("You are a Santa Claus. Buy presents and deliver them to the children.")

        
def run_agent(user_prompt, history, invariant_api_key):
    prompt = "Deliver Xbox to John."
    messages, gradio_messages = agent.run_santa_agent(prompt)
    # messages = [
    #     {"role": "user", "content": "hi there"},
    #     {"role": "assistant", "content": "bye bye"},
    # ]

    agent_params = {"system_prompt": user_prompt}

    # run command invariant test test_agent.py --agent-params '{"system_prompt": "you are santa"}'
    import subprocess
    out = subprocess.run([
        "INVARIANT_API_KEY=" + invariant_api_key,
        "invariant", "test", "test_agent.py", 
        "--agent-params", json.dumps(agent_params),
        "--push", "--dataset_name", "santa_agent",
    ], capture_output=True, text=True)
    print(out.stdout)
    print(out.stderr)
        
    return gradio_messages, "", out.stdout



with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(
                type="messages",
                label="Santa Agent", 
                avatar_images=[
                    None,
                    "https://invariantlabs.ai/theme/images/logo.svg"
                ],
            )
        with gr.Column(scale=1):
            console = gr.TextArea(label="Console Output", interactive=False)
    input = gr.Textbox(lines=1, label="System Prompt")
    invariant_api_key = gr.Textbox(lines=1, label="Invariant API Key")
    input.submit(run_agent, [input, chatbot, invariant_api_key], [chatbot, input, console])
    input.submit(lambda: gr.update(visible=False), None, [input])

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