import json
import os
import re
import invariant.testing.functional as F
import gradio as gr
from agent import SantaAgent

INITIAL_SYTSTEM_PROMPT = "You are a Santa Claus. Buy presents and deliver them to the children."
INITIAL_CHABOT = [
    {"role": "user", "content": "Could you please deliver Xbox to John?"},
]
INITIAL_STATE = ""

agent = SantaAgent(INITIAL_SYTSTEM_PROMPT)

def run_testing(agent_params, invariant_api_key):
    env={
        "INVARIANT_API_KEY": invariant_api_key,
        "OPENAI_API_KEY": os.environ["OPENAI_API_KEY"],
        "PATH": os.environ["PATH"]
    }
    import subprocess
    out = subprocess.run([
        "invariant", "test", "test_agent.py", 
        "--agent-params", json.dumps(agent_params),
        "--push", "--dataset_name", "santa_agent",
    ], capture_output=True, text=True, env=env)
    url = re.search(r"https://explorer.invariantlabs.ai/[\-_a-zA-Z0-9/]+", out.stdout).group(0)
    return url

with gr.Blocks() as demo:
    # Add state at the beginning of the Blocks
    results_state = gr.State(INITIAL_STATE)
    
    gr.Markdown("""
    ## Prompt the Santa Agent
    * Find a system prompt that passes all the tests
    """)
    input = gr.Textbox(lines=1, label="""System Prompt""", value=INITIAL_SYTSTEM_PROMPT)
    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(
                type="messages",
                label="Example interaction",
                value=INITIAL_CHABOT,
                avatar_images=[
                    None,
                    "https://invariantlabs.ai/theme/images/logo.svg"
                ],
            )
        with gr.Column(scale=1):
            console = gr.Button(value="Console Output", visible=False)

    invariant_api_key = gr.Textbox(lines=1, label="""Invariant API Key - you can play without it, but to obtain full score please register and get the key at https://explorer.invariantlabs.ai/settings""")

    def run_agent_with_state(user_prompt, history, invariant_api_key, state, is_example=False):
        # messages, gradio_messages = agent.run_santa_agent(prompt)
        gradio_messages = [
            {"role": "user", "content": "Could you please deliver Xbox to John?"},
            {"role": "assistant", "content": "I'm sorry, but I can't deliver presents. I'm just a chatbot."},
        ]

        if not invariant_api_key.startswith("inv"):
            return gradio_messages, "", "Please enter a valid Invariant API key to get the score!", state

        agent_params = {"system_prompt": user_prompt}
        
        return gradio_messages, "", "Testing in progress...", [agent_params, invariant_api_key]

    def update_console(state):
        if type(state) == list:
            agent_params, invariant_api_key = state[0], state[1]
            return gr.update(value="Testing in progress...", interactive=False, visible=True), (agent_params, invariant_api_key)
        if type(state) == tuple:
            agent_params, invariant_api_key = state
            url = run_testing(agent_params, invariant_api_key)
            return gr.update(value="Open results", link=url, visible=True, interactive=True), url
        if type(state) == str and state.startswith("https"):
            return gr.update(value="Open results", link=state, visible=True, interactive=True), state
        return gr.update(value="Testing in progress..."), state

    input.submit(run_agent_with_state, [input, chatbot, invariant_api_key, results_state], 
                [chatbot, input, console, results_state])
    input.submit(lambda: gr.update(visible=False), None, [input])

    submit = gr.Button("Submit")
    submit.click(run_agent_with_state, [input, chatbot, invariant_api_key, results_state], 
                [chatbot, input, console, results_state])
    submit.click(lambda: gr.update(visible=False), None, [input])

    reset = gr.Button("Reset")
    def reset_state():
        return (
            gr.update(value=INITIAL_SYTSTEM_PROMPT, visible=True),  # input
            INITIAL_CHABOT, # chatbot
            INITIAL_STATE, # results_state
            gr.update(visible=False), # console
        )
    
    reset.click(reset_state, None, [input, chatbot, results_state, console])

    timer = gr.Timer(value=1.0)
    timer.tick(update_console, results_state, [console, results_state])

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