"""This file contains various formatting checks used to reprompt an agent for correctly formatted responses."""

import ast
from mm_agents.vlaa_gui.utils.common_utils import (
    extract_agent_functions,
    parse_code_from_string,
    create_pyautogui_code,
    split_thinking_response,
)

single_action_check = lambda response: (
    len(extract_agent_functions(parse_code_from_string(response))) == 1
)
single_action_error_msg = (
    "Incorrect code: There must be a single agent action in the code response."
)
SINGLE_ACTION_FORMATTER = lambda response: (
    single_action_check(response),
    single_action_error_msg,
)


def _attempt_code_creation(agent, code, obs):
    """Attempts to create a pyautogui code snippet from the response code"""
    try:
        return create_pyautogui_code(agent, code, obs)
    except Exception as e:
        return None


_CODE_VALIDATION_BYPASS_ACTIONS = {
    # These actions have side-effects at eval-time (e.g. they can trigger long-running
    # work out-of-band). The formatter should not execute them while validating.
    "call_search_agent",
    "call_code_agent",
    # `agent.click(...)` is safe to bypass because it is validated elsewhere and the
    # grounding step will still validate/execute it.
    "click",
}


def _is_single_agent_call(code: str, allowed_action_names: set[str]) -> bool:
    """Returns True if the code is exactly a single `agent.<action>(...)` call."""
    try:
        tree = ast.parse(code)
    except SyntaxError:
        return False

    if len(tree.body) != 1:
        return False

    stmt = tree.body[0]
    if not isinstance(stmt, ast.Expr):
        return False
    call = stmt.value
    if not isinstance(call, ast.Call):
        return False
    func = call.func
    if not isinstance(func, ast.Attribute):
        return False
    if func.attr not in allowed_action_names:
        return False
    if not isinstance(func.value, ast.Name):
        return False
    return func.value.id == "agent"


def code_valid_check(agent, obs, response):
    code = parse_code_from_string(response)

    # Some actions must bypass code validation to avoid eval-time side effects.
    if _is_single_agent_call(code, _CODE_VALIDATION_BYPASS_ACTIONS):
        return True

    return _attempt_code_creation(agent, code, obs) is not None


code_valid_error_msg = "Incorrect code: The agent action must be a valid function and use valid parameters from the docstring list."
CODE_VALID_FORMATTER = lambda agent, obs, response: (
    code_valid_check(agent, obs, response),
    code_valid_error_msg,
)

thoughts_answer_tag_check = lambda response: split_thinking_response(response)[1] != ""
thoughts_answer_tag_error_msg = "Incorrect response: The response must contain both <thoughts>...</thoughts> and <answer>...</answer> tags."
THOUGHTS_ANSWER_TAG_FORMATTER = lambda response: (
    thoughts_answer_tag_check(response),
    thoughts_answer_tag_error_msg,
)

integer_answer_check = lambda response: (
    split_thinking_response(response)[0].strip().isdigit()
)
integer_answer_error_msg = (
    "Incorrect response: The <answer>...</answer> tag must contain a single integer."
)
INTEGER_ANSWER_FORMATTER = lambda response: (
    integer_answer_check(response),
    integer_answer_error_msg,
)
