from smolagents import (
    DuckDuckGoSearchTool,
    VisitWebpageTool,
    ToolCallingAgent,
    CodeAgent,
)

from src.models import OllamaModel
from src.tools_registry import ToolsRegistry


def check_reasoning_and_answer(final_answer, agent_memory):
    multimodal_model = OllamaModel(model_id="qwen3:14b", tools_registry=ToolsRegistry)
    prompt = (
        "You are a LLM which task is to check the summary of a reasoning and an answer given by other managed LLMs when queried to realise a task."
        "Here is a user-given task and the agent steps: "
        f"{agent_memory.get_succinct_steps()}."
        "Here is the final answer:"
        f"{final_answer}"
        "Please check that the reasoning process and the answer are correct: do they correctly answer the given task?"
        "To help you realize this task, here is some guidances:"
        "1. First, list reasons why yes/no"
        "2. Then check the answer is either a string (one or a few words), a number, or a comma separated list of strings or floats, unless specified otherwise. There is only one correct answer.  "
        "3. Conclude writing your final decision: PASS in caps lock if it is satisfactory, FAIL if it is not."
        "4. Don't be harsh, if the task couldn't achieve a perfect reasoning because direct data couldn't be found, just PASS in caps lock."
    )
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": prompt,
                }
            ],
        }
    ]
    output = multimodal_model(messages).content
    print("Feedback: ", output)

    if "FAIL" in output:
        raise Exception(output)

    return True


web_agent = ToolCallingAgent(
    tools=[
        DuckDuckGoSearchTool(),
        VisitWebpageTool(),
    ],
    model=OllamaModel(model_id="qwen3:14b", tools_registry=ToolsRegistry),
    name="web_agent",
    description="Browses the web to find information",
    verbosity_level=0,
    max_steps=10,
)


manager_agent = CodeAgent(
    model=OllamaModel(model_id="qwen3:14b"),
    tools=[],
    managed_agents=[web_agent],
    planning_interval=5,
    verbosity_level=2,
    final_answer_checks=[check_reasoning_and_answer],
    max_steps=15,
)