import os
import gradio as gr
from transformers import ReactCodeAgent, HfEngine, Tool

from gradio_agentchatbot import (
    AgentChatbot,
    stream_from_transformers_agent,
    ChatMessage,
)
from dotenv import load_dotenv
from huggingface_hub import login
from transformers.agents.default_tools import (
    BASE_PYTHON_TOOLS,
    LIST_SAFE_MODULES,
    evaluate_python_code,
)

# to load SerpAPI key
load_dotenv()
login(os.getenv("HUGGINGFACEHUB_API_TOKEN"))


llm_engine = HfEngine(model="meta-llama/Meta-Llama-3-70B-Instruct")

authorized_imports = ["numpy"]

agent = ReactCodeAgent(
    llm_engine=llm_engine,
    tools=[],
    additional_authorized_imports=authorized_imports,
    max_iterations=10,
)

class FinalAnswerToolWithVerification(Tool):
    name = "final_answer"
    description = "Provides a final answer to the given problem"
    inputs = {
        "answer": {"type": "text", "description": "The final answer to the problem"}
    }
    output_type = "any"

    def forward(self, answer):
        if "def test" not in answer:
            raise Exception(
                "I can only accept from you a code snippet answer that defines test functions in python, anything else will not work. PLEASE PROVIDE ME A FULL CODE SNIPPET CONTAINING THE DEFINITION OF THE TESTS."
            )
        return answer

final_answer_tool = FinalAnswerToolWithVerification()

agent._toolbox.update_tool(final_answer_tool)

function = """import numpy as np
def moving_average(x, w):
    return np.convolve(x, np.ones(w), 'valid') / w"""

task = """I will give you a basic function that I've created.
Now I want you to generate a set of unit tests for these functions, check that they run, and give them to me.
Please follow these steps in order:
1. Define and run the function given o you, so that it gets defined in your interpreter.
2. Generate one test function as a python blob, with assert statements
3. Run the test function in a code snippet and make sure the tests pass
4. Return to me the complete TEST function (not the original function) as a string code snippet.
---
Example:
Here is your function:
```py
def get_even_numbers(numbers):
    even_numbers = []
    for number in numbers:
        if number % 2 == 0:
            even_numbers.append(number)
    return even_numbers
```
Now generate test functions for me!
Thought: Let's re-define the given function and generate a test.
Code:
```py
def get_even_numbers(numbers):
    even_numbers = []
    for number in numbers:
        if number % 2 == 0:
            even_numbers.append(number)
    return even_numbers
def test_get_even_numbers():
    assert get_even_numbers([1, 2, 3, 4, 5]) == [2, 4]
    print("No error found!")
test_get_even_numbers()
```
Observation: "No error found!"
Thought: the interpreter ran tests with no error. So we can return the function IN A TEXT SNIPPET.
Code:
```py
fianl_answer_snippet = \"\"\"
def test_get_even_numbers():
    assert get_even_numbers([1, 2, 3, 4, 5]) == [2, 4]
    print("No error found!")
\"\"\"
final_answer(final_answer_snippet)
```
---
Now proceed!
Here is your function:
```py
<<function>>
```
Now generate test functions for me!
"""


def interact_with_agent(prompt):
    full_prompt = task.replace("<<function>>", prompt)
    messages = []
    messages.append(ChatMessage(role="user", content=full_prompt))
    yield messages
    for msg in stream_from_transformers_agent(agent, full_prompt):
        messages.append(msg)
        yield messages
    yield messages


with gr.Blocks(theme="soft") as demo:
    gr.Markdown("""### Python test generator
Write your function in the left textbox, and the agent on the right will generate tests for you!""")
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(
                lines=1, label="Your function to test", value=function
            )
            submit = gr.Button("Generate tests!")
        with gr.Column():
            chatbot = AgentChatbot(label="Agent")

    submit.click(interact_with_agent, [text_input], [chatbot])

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