import copy
import sys
from collections import deque

from constants.constants import TRAJECTORY_DIR_TEMPLATE
from utils.file_utils import resolve_path_conflict

DEBUG_ENABLED = sys.gettrace() is not None
if DEBUG_ENABLED:
    print("\n====================\nINTERACTIVE DEBUG MODE ACTIVATED\n====================\n")


import argparse
import json
import logging
import os
import random
import re
import shutil
import tempfile
import time
import traceback
import warnings
from datetime import datetime
from functools import partial
from pathlib import Path
from typing import Any

from beartype.roar import BeartypeDecorHintPep585DeprecationWarning

from agent.modular_agent import ModularAgent, construct_modular_agent
from agent.teacher_forcing import TeacherForcingAgent
from browser_env import ActionTypes, ScriptBrowserEnv, StateInfo, Trajectory, create_stop_action
from browser_env.helper_functions import RenderHelper, get_action_description
from browser_env.utils import get_task_id_from_config_file, load_task_config

# Constants
from constants import (
    AGENTS_CONFIG_DIR,
    AUTH_DIR,
    DEFAULT_RESULTS_DIR,
    HTMLS_SUBDIR,
    LM_LOGS_DIR_TEMPLATE,
    RESULT_DIR_TEMPLATE,
)
from evaluation_harness import evaluator_router
from llms.setup_utils import restore_api_keys_to_file
from utils import timing_utils as timer

# Utilities
from utils.data_recorder_reflection import DataRecorder
from utils.eval_utils import get_agent_config, set_seed
from utils.logger_utils import logger, save_log_file, save_log_file_path
from utils.signal_utils import signal_manager
from utils.string_utils import safe_format
from utils.timing_utils import dump_timings, set_timings_global_id
from utils.trajectory_view import TrajectoryView
from vwa_utils.captioner_utils import define_captioning_fn
from vwa_utils.vwa_utils import (
    early_stop,
    get_config_base_dir_from_txt,
    get_domain_from_test_config_dir,
    get_tasks_with_trajectory,
    maybe_reset_environments,
    reset_cookies_with_retry,
)

# Filter warnings
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
warnings.filterwarnings("ignore", category=BeartypeDecorHintPep585DeprecationWarning)
warnings.filterwarnings(
    "ignore", category=FutureWarning, message="`resume_download` is deprecated and will be removed in version 1.0.0."
)
warnings.filterwarnings(
    "ignore", category=UserWarning, message='Field "model_id" has conflict with protected namespace "model_".'
)

# Handler for termination signals
signal_manager.add_cleanup_function(restore_api_keys_to_file)

from trajectory import Trajectory as GeneralTrajectory
from trajectory import VWA_Trajectory


def set_debug_inputs(args: argparse.Namespace) -> None:
    # ANCHOR @debug interactive
    global DEBUG_ENABLED
    if not DEBUG_ENABLED:
        return

    # args.task_list = "rerun_reddit.txt"

    args.manual_input = False
    args.test_start_idx = 174  # inclusive
    args.test_end_idx = 174  # inclusive
    domain = "classifieds"

    # args.csv_previous_executions = "consolidated_csvs/consolidated_summary.csv"

    args.sleep_after_execution = 0.0
    args.log_html = True
    args.repeating_action_failure_th = 10
    args.log_trajectory = True
    args.show_scroll_bar = True
    args.require_all_sites_up = 1
    args.docker_instance_id = 10
    args.force_reset = False
    args.max_steps = 30
    args.skip_env_reset = True
    # args.task_list = "tasks.txt"
    args.test_config_base_dir = f"config_files/vwa_not_vague/test_{domain}"

    args.agent_config_file = "reflexion_2p-verifier-oracle-execution.yaml"  # "verifier_t-30_yes-expert_verifier_dcv5_r4_inf_q_fixK_nocap.yaml"  #  "noverifier_t-30.yaml"  # "noverifier_t-30.yaml"  # "debug.yaml"
    args.render = False

    args.agent_captioning_model_device = "server-cuda"  # 'cuda'
    # args.eval_captioning_model_device = "server-cuda"  # 'cuda'
    args.no_caption_text_obs = False

    args.viewport_width = 1280  # Default: 1280
    args.viewport_height = 2048  # Default: 720 for small context window models | 2048 for large context window models

    args.fuzzy_match_provider = "google"
    args.render_screenshot = True
    args.observation_type = "image_som"  # accessibility_tree, image_som, accessibility_tree_with_captioner, html, image

    # less used
    args.path_raw_prompts = "agent/prompts/raw/base"
    args.result_dir = "results/debug"
    args.local = False
    args.save_images = False
    args.deployment_mode = "automodel"
    args.slow_mo = 0
    args.tgi_model_endpoint = "http://127.0.0.1:8080"

    if "debug" in args.result_dir and os.path.exists(os.path.abspath(args.result_dir)):
        shutil.rmtree(os.path.abspath(args.result_dir))
        os.makedirs(os.path.abspath(args.result_dir))


# ===============================================================================
# CMD arguments
# ===============================================================================


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run end-to-end evaluation on the benchmark")

    # --- Agent configuration ---
    parser.add_argument(
        "--agent_config_file",
        type=str,
        default="agent_config_base.yaml",
        help="Filename for the YAML config file.",
    )

    parser.add_argument(
        "--active_modules",
        type=str,
        default="",
        help=(
            "If provided, activates certain modules defined in `agent_config_file`. "
            "Leave empty to activate all (default)."
            "Use <executor>:<text_refiner> to activate all modules containing the strings 'executor' and 'text_refiner'"
        ),
    )

    # --- Prompt configuration ---
    parser.add_argument(
        "--path_raw_prompts",
        type=str,
        default="./agent/prompts/raw/base",
        help="Path to raw files to build json prompts.",
    )

    # --- Subdir to store outputs ---
    parser.add_argument("--result_dir", type=str, nargs="?", const="", default="")

    # --- Execution configuration ---
    parser.add_argument("--manual_input", action="store_true", help="Low level action input comes from the user.")

    parser.add_argument(
        "--deployment_mode",
        choices=["tgi", "vllm", "automodel"],
        default="automodel",
        help="Deployment mode for the hugging-face models. If TGI, tgi_model_endpoint should be provided.",
    )

    parser.add_argument("--sleep_after_execution", type=float, default=0.0)

    parser.add_argument(
        "--tgi_model_endpoint",
        type=str,
        default="http://127.0.0.1:8080",  # if using TGI, must deploy in the address provided here
        help="Endpoint where the model is being deployed. Defaults to localhost:8080.",
    )

    parser.add_argument(
        "--docker_instance_id",
        type=str,
        default="0",
        help="Instance ID for homepage websites and env containers.",
    )

    # TGI execution
    parser.add_argument(
        "--local",
        action="store_true",
        help="Deploy a local TGI server running on --tgi_model_endpoint. Requires text-generation-launcher installation.",
    )

    # vllm execution
    parser.add_argument("--num_gpus", type=int, default=1)

    parser.add_argument(
        "--max_model_len", type=int, default=-1, help="Model context length in vLLM. Use if needs to reduce GPU memory."
    )

    parser.add_argument("--flash_attn", action="store_true", help="Uses flash attention in AutoModel engine.")

    parser.add_argument(
        "--eager", action="store_true", help="Force eager mode if using vLLM. Uses less GPU memory, but less efficient."
    )

    parser.add_argument("--debugging", action="store_true", help="Set debugging inputs.")

    parser.add_argument("--trajectory_html_path", type=str, nargs="?", const=None, default=None)

    # --- Task configuration ---
    parser.add_argument(
        "--test_config_base_dir", type=str, default="", nargs="?", help="Path to the test configuration base directory."
    )
    parser.add_argument("--test_start_idx", type=int, default=0)
    parser.add_argument("--test_end_idx", type=int, default=10000)
    parser.add_argument("--task_list", nargs="?", const=None, type=str, default=None)
    parser.add_argument("--max_tasks", type=int, default=None, help="Maximum number of tasks to evaluate.")
    parser.add_argument("--shuffle_tasks", action="store_true", help="Shuffle the task list.")
    parser.add_argument("--seed", type=int, default=42, help="Set seed for testing.")

    # ===========================================================================
    # (V)WA specific
    # ===========================================================================
    parser.add_argument("--env", type=str, default="vwa", help="Environment to run the test on.", choices=["vwa", "wa"])
    parser.add_argument("--force_reset", action="store_true", help="Reset websites for each task before its execution.")
    # --- Captioning configuration ---
    parser.add_argument(
        "--agent_captioning_model_device",
        type=str,
        default="cuda",
        help="Device to run captioning model on. By default, runs it on CUDA.",
    )
    parser.add_argument(
        "--skip_env_reset",
        action="store_true",
        help="Skip resetting the environment for each task. Useful for debugging.",
    )

    parser.add_argument(
        "--eval_captioning_model_device",
        type=str,
        default="cpu",
        help="Device to run eval captioning model on. By default, runs it on CPU.",
    )

    parser.add_argument(
        "--eval_captioning_model",
        type=str,
        default="Salesforce/blip2-flan-t5-xl",
        choices=["Salesforce/blip2-flan-t5-xl", "llava-hf/llava-1.5-7b-hf"],
        help="Captioning backbone for VQA-type evals.",
    )

    parser.add_argument(
        "--no_caption_text_obs",
        action="store_true",
        help="If true, does not caption the text observations of the webpage, even if captioner is provided.",
    )

    # --- Observation configuration ---
    parser.add_argument(
        "--observation_type",
        choices=[
            "accessibility_tree",  # text: actree of the webpages
            "accessibility_tree_with_captioner",  # text: actree of the webpages; caption model generates description to images and add to actree
            "image_som",  # image: SOM-marked screenshot of webpage; text: ID of items marked in the screenshot
            "html",  # TODO: codebase dont have action parser for this
            "image",  # TODO: codebase dont to have action parser for this
        ],
        default="accessibility_tree",
        help="Observation type",
    )

    parser.add_argument(
        "--show_scroll_bar",
        action="store_true",
        help="Show the scroll bar in the observation",
    )

    parser.add_argument(
        "--current_viewport_only",
        action="store_true",
        help="Only use the current viewport for the observation",
    )
    parser.add_argument("--viewport_width", type=int, default=1280)
    parser.add_argument("--viewport_height", type=int, default=2048)

    # --- Evaluation configuration ---
    parser.add_argument(
        "--parsing_failure_th",
        type=int,
        default=3,
        help="When consecutive parsing failure exceeds this threshold, the agent will stop",
    )

    parser.add_argument(
        "--repeating_action_failure_th",
        type=int,
        default=30,
        help="When consecutive repeated actions exceed this threshold, the agent will terminate early.",
    )

    parser.add_argument(
        "--fuzzy_match_provider",
        type=str,
        default="google",
        choices=["openai", "google", "huggingface"],
        help="LLM provider for fuzzy matching evaluation.",
    )

    parser.add_argument(
        "--max_steps", type=int, default=30, help="Max number of environment steps allowed. If exceeded, FAIL."
    )

    # --- Outputs/results configuration ---
    parser.add_argument("--render_screenshot", action="store_true")
    parser.add_argument("--log_html", action="store_true")

    # real-time browser rendering
    parser.add_argument("--render", action="store_true", help="Shows the browser in the screen.")
    parser.add_argument(
        "--slow_mo",
        type=int,
        default=0,
        help="Slow down the browser by the specified amount",
    )
    parser.add_argument(
        "--require_all_sites_up", type=int, choices=[0, 1], default=1, help="Require all sites to be up."
    )

    parser.add_argument("--log_trajectory", action="store_true")

    parser.add_argument("--save_trace_enabled", action="store_true")

    args = parser.parse_args()

    return args


def build_test_file_list(args: argparse.Namespace) -> list[str]:
    test_file_list = []

    if args.task_list:
        logger.info(f"\nExecuting tasks from: {args.task_list}")
        # Read task file and extract JSON paths (only for numeric lines)
        with open(args.task_list, "r") as f:
            test_file_list.extend(
                [
                    os.path.join(args.test_config_base_dir, f"{line.strip()}.json")
                    for line in f.readlines()
                    if line.strip().isdigit()
                ]
            )
        # Copy the task list file to the result directory if not already there
        src_path = os.path.abspath(args.task_list)
        dest_path = os.path.abspath(os.path.join(args.result_dir, os.path.basename(args.task_list)))
        if src_path != dest_path:
            shutil.copyfile(args.task_list, dest_path)
    else:
        logger.info(f"\nExecuting tasks in the range: {args.test_start_idx} to {args.test_end_idx}")
        test_file_list = [
            os.path.join(args.test_config_base_dir, f"{i}.json")
            for i in range(args.test_start_idx, args.test_end_idx + 1)
        ]

    # If a trajectory HTML path is provided, update the task list using the intersection with tasks that have trajectories
    if args.trajectory_html_path:
        logger.info(f"Set of tasks updated to tasks with previous trajectories in {args.trajectory_html_path}")
        test_file_list = get_tasks_with_trajectory(args.trajectory_html_path, args.test_config_base_dir, test_file_list)

    # Remove finished tasks # TODO: utilize get_unfinished from experiments_utils
    # test_file_list = get_unfinished(test_file_list, args.result_dir)

    if num_tasks := len(test_file_list) == 0:
        logger.info("No task left to run")
        sys.exit(0)

    if args.max_tasks is not None and args.max_tasks > 0 and num_tasks > args.max_tasks:
        logger.info(f"`max_tasks` set to {args.max_tasks}. Evaluating the first {args.max_tasks} tasks of {num_tasks}.")
        test_file_list = test_file_list[: args.max_tasks]

    if args.shuffle_tasks:
        logger.info(f"Shuffling tasks. Seed: {args.seed}")
        random.shuffle(test_file_list)

    return test_file_list


# Pre-process args
def preprocess_args(args: argparse.Namespace) -> None:
    # check the whether the action space is compatible with the observation space
    args.action_set_tag = args.agents_configs["executor"]["action_set_tag"]

    # removed `som_image`` because `id_accessibility_tree` with `som_image` gives wrong hint in get_action_description(..) (see the function for more details)
    if args.action_set_tag == "id_accessibility_tree" and args.observation_type not in [
        "image",
        "accessibility_tree",
        "accessibility_tree_with_captioner",
    ]:
        raise ValueError(
            f"Action type {args.action_set_tag} is incompatible with the observation type {args.observation_type}"
        )

    if args.action_set_tag == "som" and args.observation_type not in ["image_som"]:  # image is not supported yet
        raise ValueError(
            f"Action type {args.action_set_tag} is incompatible with the observation type {args.observation_type}"
        )

    if args.observation_type in ["accessibility_tree_with_captioner"] and args.captioning_model is None:
        args.captioning_model = "Salesforce/blip2-flan-t5-xl"
        logging.warning(
            f"Observation type `{args.observation_type}` requires captioner but none provided. \
            captioning model set to {args.captioning_model} (default)"
        )

    # Create json task files (0.json, 1.json, ...)
    # generate_test_data.main() # Commented to avoid race condition

    # If trajectory_html_path provided, verify it exists
    if args.trajectory_html_path:
        if not os.path.exists(args.trajectory_html_path):
            raise ValueError(f"Trajectory HTML path {args.trajectory_html_path} does not exist.")

    # If task_list provided, get the test_config_base_dir if available
    if args.task_list:
        test_config_base_dir = get_config_base_dir_from_txt(args.task_list)
        if test_config_base_dir:
            args.test_config_base_dir = test_config_base_dir

    # Verify existent tasks in config folder and regularize task indices
    if args.task_list is None:
        test_files = os.listdir(args.test_config_base_dir)
        test_files = [f for f in test_files if f.endswith(".json")]
        if len(test_files) == 0:
            raise ValueError(f"No tasks found in {args.test_config_base_dir}")
        task_ids = [int(os.path.basename(f).split(".")[0]) for f in test_files]
        min_idx, max_idx = min(task_ids), max(task_ids)
        args.test_start_idx = min_idx if args.test_start_idx < min_idx else args.test_start_idx
        args.test_end_idx = max_idx if args.test_end_idx > max_idx else args.test_end_idx

    # ---------------------------------------------------------------------------
    # Results and log directory paths
    # ---------------------------------------------------------------------------

    # --- Construct result directory path ----
    # Creates a default `args.result_dir` if not provided
    if not args.result_dir:
        top_dir = DEFAULT_RESULTS_DIR

        # Extract the domain name to annotate the result directory
        args.domain = domain = get_domain_from_test_config_dir(args.test_config_base_dir)

        # Annotate model and current date and time
        date_annotate = datetime.now().strftime("%Y-%m-%d-%H%M")
        model = args.agents_configs["executor"]["lm_config"]["model"]  # TODO: change if using multiple models
        args.result_dir = safe_format(
            RESULT_DIR_TEMPLATE, results_dir=top_dir, model=model, domain=domain, annotation=date_annotate
        )
        # Example: results/gpt-4o-mini-2024-07-18/reddit_2025-02-22-1430

    # NOTE: This is helpful if parallel runs; runs will all be saved at the initial `args.result_dir`, but separated by their corresponding tasks.
    # Create a subdir in `args.result_dir` to hold results for a range of tasks or a task list
    if args.task_list is None:
        args.result_dir = f"{args.result_dir}/{args.test_start_idx}-{args.test_end_idx}"
    else:
        args.result_dir = f"{args.result_dir}/{Path(args.task_list).stem}"
    args.result_dir = str(resolve_path_conflict(args.result_dir, int_suffix=True))

    # Create the result directory, if not exists
    os.makedirs(args.result_dir, exist_ok=True)

    # Create the traces directory, if not exists
    if args.save_trace_enabled:
        os.makedirs(os.path.join(args.result_dir, "traces"), exist_ok=True)

    # Save logger files paths to a text file in results folder
    save_log_file_path(args.result_dir)

    # Base name of files to save execution HTML traces and LM call logs.
    args.htmls_dir = os.path.join(args.result_dir, HTMLS_SUBDIR)
    if args.log_html:
        os.makedirs(args.htmls_dir, exist_ok=True)
    args.lm_logs_dir = LM_LOGS_DIR_TEMPLATE.format(result_dir=args.result_dir)
    os.makedirs(args.lm_logs_dir, exist_ok=True)


# ===============================================================================
# Evaluation
# ===============================================================================


def error_handler(error: Exception, result_dir: str, data_recorder: DataRecorder, raw_config_file: str) -> bool:
    logger.error(repr(error), exc_info=True)
    with open(Path(result_dir) / "error.txt", "a") as f:
        # Dump config file content
        if raw_config_file:
            f.write(f"[Config file]: {raw_config_file}\n")

        f.write(f"[Error trace] {repr(error)}\n")
        f.write(traceback.format_exc())  # write stack trace to file
    data_recorder.num_failed_executions += 1
    return False


def init_environment(args: argparse.Namespace) -> ScriptBrowserEnv:
    captioning_fn = None if args.no_caption_text_obs else args.caption_img_fn
    return ScriptBrowserEnv(
        headless=not args.render,
        slow_mo=args.slow_mo,
        observation_type=args.observation_type,
        current_viewport_only=args.current_viewport_only,
        viewport_size={
            "width": args.viewport_width,
            "height": args.viewport_height,
        },
        save_trace_enabled=args.save_trace_enabled,
        sleep_after_execution=args.sleep_after_execution,
        captioning_fn=captioning_fn,
        show_scroll_bar=args.show_scroll_bar,
    )


# ANCHOR: Evaluation
# TODO: clean logic for config_file; modularize test loop
def test(args: argparse.Namespace, agent: ModularAgent, config_file_list: list[str]) -> None:
    test_start_time = time.time()  # TODO: remove this
    domain = get_domain_from_test_config_dir(args.test_config_base_dir)
    args.domain = domain

    # Initialize data recorder to record experiment data
    data_recorder = DataRecorder(args.result_dir, config_file_list, args.test_config_base_dir)

    # Set early stop thresholds
    max_steps = args.max_steps
    early_stop_thresholds = {
        "parsing_failure": args.parsing_failure_th,
        "repeating_action": args.repeating_action_failure_th,
    }

    # Initialize environment
    env = init_environment(args)

    # ANCHOR: Iterate each task (defined in `config_file` jsons)
    config_file_queue = deque(config_file_list)
    args.attempts_per_task = {}  # task_id: num_attempts
    iteration_count = 0
    while config_file_queue:
        iteration_count += 1
        raw_config_file = config_file_queue.popleft()
        traj_logger = None
        task_id = get_task_id_from_config_file(raw_config_file)
        args.attempts_per_task[task_id] = args.attempts_per_task.get(task_id, -1) + 1
        logger.info(f"[Config file]: {raw_config_file}, attempt: {args.attempts_per_task[task_id]}")

        try:
            test_execution_success = True  # Track if any errors during execution of current task
            # Reset for classifieds is lightweight when using classifieds reset token. Forcing it.
            if domain == "classifieds":
                args.force_reset = True

            if not args.skip_env_reset:
                reset_success, sites_failed_reset = maybe_reset_environments(
                    instance_id=args.docker_instance_id,
                    env=args.env,
                    config_dict=json.loads(Path(raw_config_file).read_text()),
                    require_all_sites_up=args.require_all_sites_up,
                    force_reset=args.force_reset if iteration_count > 1 else False,
                    force_reset_site_list=[domain],
                )
                if not reset_success:
                    logger.error(
                        f"Failed to reset environments for {sites_failed_reset}, config_file: {raw_config_file}."
                    )
                    raise Exception(
                        f"Failed to reset environments for {sites_failed_reset}, config_file: {raw_config_file}"
                    )

            # Load task config
            config_dict, intent, task_id, intent_images, image_paths = load_task_config(
                raw_config_file, args.docker_instance_id
            )

            # # Reset cookies for current task
            reset_cookies_success, msg = reset_cookies_with_retry(
                docker_instance_id=args.docker_instance_id,
                sites=config_dict["sites"],
                expired_only=False,
                exc_comb=False,
                wait_for_cookies_reset=True,
                auth_folder=AUTH_DIR,
            )
            if not reset_cookies_success:
                logger.error(f"Failed to reset cookies: {msg}. Config file: {raw_config_file}", exc_info=True)
                raise Exception(f"Failed to reset cookies: {msg}. Config file: {raw_config_file}")

            # If reflexion enabled, populate reflection memory and update attempt count
            if agent.reflexion is not None:
                args.attempts_per_task[task_id] = agent.reflexion.get_num_attempts_per_task(f"{domain}_{task_id}")
                if args.attempts_per_task[task_id] > agent.reflexion.max_reflexion_attempts:
                    logger.info(f"Max reflexion attempts reached for {domain}_{task_id}: {args.max_reflexion_attempts}")
                    continue

            # Create temporary config file from config_dict
            # TODO: clean this up. For backward compatibility, still creates a file.
            with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as temp_file:
                json.dump(config_dict, temp_file)
                temp_config_file = temp_file.name

            # Time logging
            set_timings_global_id(f"{domain}-{task_id}_{args.attempts_per_task[task_id]}")
            timer.start("RUN:test")

            evaluator = evaluator_router(
                temp_config_file,
                captioning_fn=args.eval_caption_img_fn,
                fuzzy_match_prov=args.fuzzy_match_provider,
            )  # type: ignore

            # Log info #TODO: move this up once config_file logic is cleaned up
            logger.info(f"[Intent]: {intent}")

            # Data logging
            if args.log_html:
                render_helper = RenderHelper(
                    temp_config_file,
                    args.htmls_dir,
                    args.action_set_tag,
                    postfix=args.attempts_per_task[task_id],
                )

            task_start_time = time.time()
            if args.log_trajectory:
                traj_logger = VWA_Trajectory(
                    task_id=task_id,
                    objective_text=intent,
                    objective_imgs=image_paths,
                    domain=domain,
                    config_file=raw_config_file,
                    args=f"{args.result_dir}/args.json",
                )

            # Prepare agent and environment
            agent.reset(config_dict, dict(vars(args)))
            trajectory: Trajectory = []
            trajectory_view = TrajectoryView(trajectory)
            obs, info = env.reset(options={"config_file": temp_config_file})  # type: ignore
            state_info: StateInfo = {"observation": obs, "info": info}
            trajectory.append(state_info)
            meta_data: dict[str, Any] = {
                "action_str_history": ["None"],
                "task_id": task_id,
                "data_recorder": data_recorder,
                "manual_input": args.manual_input,
                "env": env,
                "trajectory": trajectory_view,
                "args": args,
                "config_file": temp_config_file,
                "evaluator": partial(evaluator, config_file=temp_config_file),  # type: ignore
                "domain_task_id": f"{domain}_{task_id}",
                "attempt_num": args.attempts_per_task[task_id],
            }
            early_stop_flag = False
            data_recorder.initialize_task(task_id, config_dict["sites"], attempt_id=args.attempts_per_task[task_id])

            # ANCHOR: Loop: execute current task
            while True:
                save_log_file(args.result_dir)  # Write log content to `args.result_dir`
                test_execution_success = True  # Track if any errors during execution of current task
                try:
                    # Check if stop threshold reached and stop
                    early_stop_flag, stop_info = early_stop(trajectory, max_steps, early_stop_thresholds)
                    if early_stop_flag:
                        action = create_stop_action(stop_info)
                        logger.info(f"Early stop: {stop_info}")
                        if args.log_trajectory and traj_logger is not None:
                            traj_logger.info.update({"early_stop": stop_info})

                    # Else, Get action on environment
                    else:
                        action = agent.next_action(
                            trajectory_view,
                            intent,
                            intent_images=intent_images,  # type: ignore
                            meta_data=meta_data,
                        )

                        if "early_stop" in action:
                            early_stop_reason = action["early_stop"]  # type: ignore
                            logger.info(f"Early stop: {early_stop_reason}")
                            action = create_stop_action(action["raw_prediction"])
                            action["early_stop"] = early_stop_reason  # type: ignore
                            early_stop_flag = True

                except ValueError as e:
                    action = create_stop_action(f"ERROR: {str(e)}")
                    logger.error(repr(e), exc_info=True)
                    if args.log_trajectory and traj_logger:
                        traj_logger.info.update({"error": str(e)})

                # Append action to trajectory
                trajectory.append(action)

                # Convert action to string to append to history to serve as input in next generations
                action_str = get_action_description(
                    action=action,
                    observation_metadata=state_info["info"]["observation_metadata"],
                    action_set_tag=args.action_set_tag,
                    prompt_constructor=agent.get_prompt_constructor(),
                )
                action_str = re.sub(r"\s+", " ", action_str.replace("\n", "")).strip()
                meta_data["action_str_history"].append(action_str)
                action["action_str_env_parsed"] = action_str  # type: ignore

                if args.log_html:
                    render_helper.render(action, state_info, meta_data, args.render_screenshot)

                if args.log_trajectory and traj_logger is not None:
                    # Update state for current step
                    traj_logger.add_step(state=state_info, action=None, metadata=meta_data)

                dump_timings(args.result_dir)

                # If STOP action, end task loop.
                if action["action_type"] == ActionTypes.STOP:
                    break

                # ANCHOR Act on environment
                obs, _, terminated, _, info = env.step(action)

                if args.log_trajectory and traj_logger is not None:
                    # Update action for current step (done after env.step because action object is modified)
                    traj_logger.update_step(step_index=-1, state=None, action=action, metadata=meta_data)

                # Update state
                state_info = {"observation": obs, "info": info}
                trajectory.append(state_info)

                if terminated:
                    trajectory.append(create_stop_action(""))  # Add action placeholder
                    break

                if str(task_id) == "319":
                    raise Exception("test")
            # END OF TASK

            # If reflexion, deep copy trajectory for verification and reflection generation (copy because oracle evaluator can change state)
            trajectory_copy = None
            if agent.reflexion is not None and agent.reflexion.should_retry(meta_data["domain_task_id"]):
                trajectory_copy = TrajectoryView(copy.deepcopy(trajectory))

            # Evalute and print score.
            score = evaluator(trajectory=trajectory, config_file=temp_config_file, page=env.page)  # type: ignore
            logger.info(
                f"[Result] ({'PASS' if score == 1 else 'FAIL'}) {raw_config_file}, attempt: {args.attempts_per_task[task_id]}"
            )
            elapsed_time = time.time() - task_start_time

            # If reflexion enabled, evaluate, generate reflection, and add config_file for next attempt
            v_response = v_feedback = ""
            if agent.reflexion is not None and trajectory_copy is not None:
                if agent.reflexion.get_eval_mode() == "oracle":
                    is_success = True if score == 1 else False

                elif agent.reflexion.get_eval_mode() == "model":
                    is_success, v_feedback, v_response, v_parsed_response = agent.get_verifier_response(
                        trajectory_copy, intent, intent_images, meta_data
                    )
                else:
                    raise ValueError(f"Unknown eval mode: {agent.reflexion.get_eval_mode()}")

                if not is_success:
                    agent.generate_reflection(trajectory_copy, intent, intent_images, meta_data)
                    # Retry with reflection
                    config_file_queue.appendleft(raw_config_file)

            # Save trace
            if args.save_trace_enabled:
                task_id_attempt = f"{task_id}_{args.attempts_per_task[task_id]}"
                env.save_trace(Path(args.result_dir) / "traces" / f"{task_id_attempt}.zip")

            # Record and save task stats
            data_recorder.update_save_data(
                task_id,
                score,
                elapsed_time,
                num_actions=len(trajectory[1::2]),
                attempt_id=args.attempts_per_task[task_id],
            )

            # Log trajectory data
            # TODO: add logic for terminated, truncated, early_stop
            if args.log_trajectory:
                assert traj_logger is not None  # for correct type checking msgs

                # Add action on last state and end episode
                traj_logger.update_step(step_index=-1, state=None, action=action, metadata=meta_data)
                traj_logger.end_episode(
                    reward=score, total_time_seconds=elapsed_time, episode_completed=True, truncated=early_stop_flag
                )

                if hasattr(agent, "request_refiner"):
                    obj_img_captions = [caption for _, caption in agent.request_refiner.image_captions.items()]  # type: ignore
                    traj_logger.objective.update({"obj_img_captions": obj_img_captions})  # type: ignore

                if agent.reflexion is not None:
                    reflexion_data = {}
                    reflexion_data["reflections"] = agent.reflexion.reflection_memory.get(
                        meta_data["domain_task_id"], []
                    )
                    reflexion_data["attempt_num"] = args.attempts_per_task[task_id]
                    reflexion_data["verifier_response"] = v_response
                    reflexion_data["verifier_feedback"] = v_feedback
                    reflexion_data["verifier_score"] = is_success
                    traj_logger.metadata.update({"reflexion_data": reflexion_data})  # type: ignore

                # Log trajectory
                traj_logger.dump_to_json(
                    output_dir=TRAJECTORY_DIR_TEMPLATE.format(
                        result_dir=args.result_dir, task_id=f"{task_id}_{args.attempts_per_task[task_id]}"
                    ),
                    filename=f"trajectory-{domain}-{task_id}.json",
                )

        except Exception as e:
            test_execution_success = error_handler(
                e, result_dir=args.result_dir, data_recorder=data_recorder, raw_config_file=raw_config_file
            )
            if not test_execution_success and traj_logger is not None and args.log_trajectory:
                traj_logger.info.update({"error": str(e)})
                traj_logger.episode_completed = False
                try:
                    traj_logger.dump_to_json(
                        output_dir=TRAJECTORY_DIR_TEMPLATE.format(
                            result_dir=args.result_dir, task_id=f"{task_id}_{args.attempts_per_task[task_id]}"
                        ),
                        filename=f"trajectory-{domain}-{task_id}.json",
                    )
                except Exception as e:
                    logger.error(f"Error dumping trajectory: {e}")
        finally:
            timer.end("RUN:test")
            data_recorder.update_unfinished_failed_tasks(task_id, test_execution_success)
            if args.log_html:
                render_helper.close()
            if os.path.exists(temp_config_file):
                os.unlink(temp_config_file)
        # End of one task

    # End of experiment
    env.close()
    test_end_time = time.time()
    scores = data_recorder.get_scores()
    logger.info(f"Average score: {sum(scores) / len(scores)}; {len(scores)} tasks")
    logger.info(
        f"NOTE: {len(config_file_list) - data_recorder.num_failed_executions} of {len(config_file_list)} completed without error."
    )
    logger.info(f"Total test time (min): {(test_end_time - test_start_time) / 60}")

    # Save summary execution stats
    data_recorder.save_execution_summary(test_end_time - test_start_time, agent.get_provider("executor"))


# LINK: Main
if __name__ == "__main__":
    os.environ["TOKENIZERS_PARALLELISM"] = "false"
    args = parse_args()

    # Set seed for numpy, random, torch
    set_seed(args.seed)

    # Set debugging inputs, if any
    if DEBUG_ENABLED:
        set_debug_inputs(args)

    # ------------------------------------------------------------------------------
    # Parse and validate configs, args
    # ------------------------------------------------------------------------------
    # Set Agent configurations
    if AGENTS_CONFIG_DIR in args.agent_config_file:
        args.agent_config_file = Path(args.agent_config_file).name

    args.agent_config_path = str(Path(AGENTS_CONFIG_DIR) / args.agent_config_file)
    args.agents_configs = get_agent_config(args.agent_config_path, active_modules_str=args.active_modules)

    # Validate and preprocess args
    preprocess_args(args)

    # Build test file list
    test_file_list = build_test_file_list(args)

    # ------------------------------------------------------------------------------
    # Dump test_file_list and args to results folder
    # ------------------------------------------------------------------------------
    with open(os.path.join(args.result_dir, "attempted_tasks.txt"), "w") as f:
        for test_file in sorted(test_file_list):
            f.write(f"{test_file}\n")

    with open(f"{args.result_dir}/args.json", "w") as f:
        json.dump(vars(args), f, indent=4)

    logger.info(f"Starting evaluation. Docker instance ID: {args.docker_instance_id}. Environment: {args.env}.")

    logger.info(f"\nTotal {len(test_file_list)} tasks left.\n")

    args.agent_captioning_model = None
    if "request_refiner" in args.agents_configs:
        args.agent_captioning_model = args.agents_configs["request_refiner"]["captioning_model"]
    args.caption_img_fn, args.eval_caption_img_fn = define_captioning_fn(args)

    # ------------------------------------------------------------------------------
    # Build Agent
    # ------------------------------------------------------------------------------
    if args.trajectory_html_path:
        # Get all HTMLs in the trajectory_html_path
        # agent: ModularAgent = TeacherForcingAgent(
        #     trajectory_html_path=args.trajectory_html_path, agents_configs=args.agents_configs
        # )
        pass
    else:
        if "reflexion" in args.agents_configs:
            args.agents_configs["reflexion"]["lm_logs_dir"] = args.lm_logs_dir  # TODO: implement for all agents
        agent = construct_modular_agent(args.agents_configs, args.caption_img_fn)

    # ------------------------------------------------------------------------------
    # Evaluate
    # ------------------------------------------------------------------------------
    # print full path to result_dir
    logger.info(f"\nResults will be saved in {os.path.abspath(args.result_dir)}")

    try:
        test(args, agent, test_file_list)
        save_log_file(args.result_dir)
        logger.info(f"Results in {os.path.abspath(args.result_dir)}")

    except Exception as e:
        logger.error(repr(e), exc_info=True)
        raise e

    finally:
        restore_api_keys_to_file()
