"""Script to run end-to-end evaluation on the benchmark.

"""

import argparse
import copy
import time
import datetime
import json
import logging
import os
import signal
import sys
from dataclasses import dataclass
from multiprocessing import Process, Manager, current_process, Queue

# Add project root to path for imports
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../.."))

from desktop_env.desktop_env import DesktopEnv
import lib_run_single
from typing import Any


#  Logger Configs {{{ #
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)

log_dir = "logs"
os.makedirs(log_dir, exist_ok=True)

datetime_str: str = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")

file_handler = logging.FileHandler(
    os.path.join("logs", "normal-{:}.log".format(datetime_str)), encoding="utf-8"
)
stdout_handler = logging.StreamHandler(sys.stdout)

stdout_handler.setLevel(logging.INFO)
file_handler.setLevel(logging.INFO)

formatter = logging.Formatter(
    fmt="\x1b[1;33m[%(asctime)s \x1b[31m%(levelname)s \x1b[32m%(module)s/%(lineno)d-%(processName)s\x1b[1;33m] \x1b[0m%(message)s"
)

stdout_handler.setFormatter(formatter)
file_handler.setFormatter(formatter)

stdout_handler.addFilter(logging.Filter("desktopenv"))
# file_handler.addFilter(logging.Filter("gui_agent_zoo"))

logger.addHandler(stdout_handler)
logger.addHandler(file_handler)
#  }}} Logger Configs #

logger = logging.getLogger("desktopenv.agent")

# Global variables for signal handling
active_environments = []
processes = []
is_terminating = False


if os.path.exists(".env"):
    from dotenv import load_dotenv

    load_dotenv()


def _default_model_dir_name(model_name: str) -> str:
    if not model_name:
        return ""
    return model_name.rsplit(".", 1)[-1] if "." in model_name else model_name


def _finalize_args(args: argparse.Namespace) -> argparse.Namespace:

    if not args.model_dir_name:
        args.model_dir_name = _default_model_dir_name(args.model)

    args.temperature = args.model_temperature

    if not args.model_region:
        args.model_region = args.region

    if not args.grounding_model_region:
        args.grounding_model_region = args.model_region

    if args.model_reasoning_effort is None:
        args.model_reasoning_effort = args.model_thinking_level

    if args.grounding_reasoning_effort is None:
        args.grounding_reasoning_effort = args.grounding_thinking_level

    if args.reflection_thinking is None:
        args.reflection_thinking = args.model_thinking

    if args.search_engine in {"", "None"}:
        args.search_engine = None

    if args.provider_name == "aws" and not args.client_password:
        args.client_password = "osworld-public-evaluation"

    return args


def distribute_tasks(test_all_meta: dict) -> list:
    all_tasks = []
    for domain, examples in test_all_meta.items():
        for example_id in examples:
            all_tasks.append((domain, example_id))
    return all_tasks


def process_signal_handler(signum, frame, env_idx):
    logger.info(f"Process {env_idx + 1} received signal {signum}. Shutting down...")
    local_vars = frame.f_locals
    active_environments = local_vars.get("active_environments", [])
    for env in active_environments:
        if env is not None:
            try:
                logger.info(f"Process {env_idx + 1} closing environment...")
                env.close()
                logger.info(f"Process {env_idx + 1} environment closed successfully")
            except Exception as e:
                logger.error(f"Process {env_idx + 1} error closing environment: {e}")
    logger.info(f"Process {env_idx + 1} shutdown complete. Exiting.")
    sys.exit(0)


def signal_handler(signum, frame):
    global is_terminating, active_environments, processes
    if is_terminating:
        return
    is_terminating = True
    logger.info(f"Received signal {signum}. Gracefully shutting down...")
    for env in active_environments:
        try:
            logger.info("Closing environment...")
            env.close()
            logger.info("Environment closed successfully")
        except Exception as e:
            logger.error(f"Error closing environment: {e}")
    for p in processes:
        if p.is_alive():
            try:
                logger.info(f"Sending termination signal to process {p.name}...")
                p.terminate()
            except Exception as e:
                logger.error(f"Error sending termination signal to process: {e}")
    time.sleep(1)
    for p in processes:
        if p.is_alive():
            try:
                logger.info(f"Forcefully terminating process {p.name}...")
                import signal as sig

                os.kill(p.pid, sig.SIGKILL)
            except Exception as e:
                logger.error(f"Error forcefully terminating process: {e}")
    logger.info("Shutdown complete. Exiting.")
    sys.exit(0)


@dataclass
class RunConfig:
    """Per-run configuration bundle for multi-run support."""

    run_idx: int
    result_dir: str
    run_label: str
    args: argparse.Namespace
    engine_params: dict
    reflection_engine_params: dict
    engine_params_for_grounding: dict
    engine_params_for_coding: dict
    engine_params_for_searcher: dict


def _build_engine_params(args: argparse.Namespace) -> tuple:
    """Build the four engine-parameter dicts from a fully-resolved *args*."""
    if args.search_engine == "None" or args.search_engine == "":
        args.search_engine = None

    engine_params = {
        "engine_type": args.model_provider,
        "model": args.model,
        "base_url": getattr(args, "model_url", ""),
        "api_key": getattr(args, "model_api_key", ""),
        "api_keys": getattr(args, "model_api_keys", []),
        "api_version": getattr(args, "api_version", ""),
        "project_id": getattr(args, "model_project_id", ""),
        "region": getattr(args, "model_region", ""),
        "aws_keys": getattr(args, "model_aws_keys", []),
        "thinking": getattr(args, "model_thinking", False),
        "thinking_budget": getattr(args, "model_thinking_budget", None),
        "thinking_level": getattr(args, "model_thinking_level", None),
        "reasoning_effort": getattr(args, "model_reasoning_effort", None),
        "temperature": getattr(args, "model_temperature", None),
        "include_thoughts": getattr(args, "model_include_thoughts", False),
    }
    reflection_engine_params = dict(engine_params)
    reflection_engine_params["thinking"] = getattr(args, "reflection_thinking", False)

    if args.endpoint_url:
        engine_params_for_grounding = {
            "engine_type": args.endpoint_provider,
            "base_url": args.endpoint_url,
            "api_key": args.endpoint_api_key,
        }
    else:
        grounding_height = args.grounding_height
        grounding_width = args.grounding_width
        if grounding_width is None:
            grounding_width = args.screen_width
        if grounding_height is None:
            grounding_height = args.screen_height * grounding_width / args.screen_width

        engine_params_for_grounding = {
            "engine_type": args.grounding_model_provider,
            "model": args.grounding_model,
            "grounding_width": grounding_width,
            "grounding_height": grounding_height,
            "thinking": getattr(args, "grounding_thinking", False),
            "thinking_budget": getattr(args, "grounding_thinking_budget", None),
            "thinking_level": getattr(args, "grounding_thinking_level", None),
            "reasoning_effort": getattr(args, "grounding_reasoning_effort", None),
            "include_thoughts": getattr(args, "grounding_include_thoughts", False),
        }
        if args.grounding_model_provider == "anthropic_bedrock":
            engine_params_for_grounding["region"] = args.grounding_model_region
            if getattr(args, "model_aws_keys", None):
                engine_params_for_grounding["aws_keys"] = args.model_aws_keys
        else:
            engine_params_for_grounding["base_url"] = args.grounding_model_url
            engine_params_for_grounding["api_key"] = args.grounding_model_api_key

    engine_params_for_coding = {
        "engine_type": args.coding_model_provider,
        "base_url": args.coding_model_url,
        "api_key": args.coding_model_api_key,
        "api_keys": getattr(args, "coding_model_api_keys", []),
        "model": args.coding_model,
        "project_id": getattr(args, "model_project_id", ""),
        "region": getattr(args, "model_region", ""),
        "aws_keys": getattr(args, "model_aws_keys", []),
    }

    engine_params_for_searcher = {
        "engine_type": args.searcher_provider,
        "type": args.searcher_type,
        "model": args.searcher_model,
        "api_key": args.searcher_api_key,
        "api_keys": getattr(args, "searcher_api_keys", []),
        "base_url": getattr(args, "searcher_url", ""),
        "budget": args.searcher_budget,
        "reasoning_effort": getattr(args, "searcher_reasoning_effort", None),
        "project_id": getattr(args, "model_project_id", ""),
        "region": getattr(args, "model_region", ""),
        "aws_keys": getattr(args, "model_aws_keys", []),
    }

    return (
        engine_params,
        reflection_engine_params,
        engine_params_for_grounding,
        engine_params_for_coding,
        engine_params_for_searcher,
    )


def build_run_configs(shared_args: argparse.Namespace) -> list:
    """Build RunConfig objects from direct CLI args or JSON run overrides."""
    raw_shared_args = copy.deepcopy(shared_args)
    _finalize_args(shared_args)

    if shared_args.runs_file:
        with open(shared_args.runs_file, "r", encoding="utf-8") as f:
            runs = json.load(f)
        if not isinstance(runs, list) or not runs:
            raise ValueError("--runs-file must contain a non-empty JSON array")

        valid_keys = set(vars(raw_shared_args))
        run_configs: list[RunConfig] = []
        for i, entry in enumerate(runs):
            if not isinstance(entry, dict):
                raise ValueError(
                    f"Each entry in --runs-file must be a JSON object (entry {i})"
                )

            unknown_keys = sorted(
                key
                for key in entry
                if key not in valid_keys and key not in {"name", "run_label"}
            )
            if unknown_keys:
                raise ValueError(
                    f"Unknown keys in --runs-file entry {i}: {', '.join(unknown_keys)}"
                )

            run_args = copy.deepcopy(raw_shared_args)
            for key, value in entry.items():
                if key in {"name", "run_label"}:
                    continue
                setattr(run_args, key, value)

            _finalize_args(run_args)
            ep, rep, epg, epc, eps = _build_engine_params(run_args)
            run_label = (
                entry.get("run_label")
                or entry.get("name")
                or entry.get("model_dir_name")
                or entry.get("result_dir")
                or f"run-{i}"
            )
            run_configs.append(
                RunConfig(
                    run_idx=i,
                    result_dir=run_args.result_dir,
                    run_label=run_label,
                    args=run_args,
                    engine_params=ep,
                    reflection_engine_params=rep,
                    engine_params_for_grounding=epg,
                    engine_params_for_coding=epc,
                    engine_params_for_searcher=eps,
                )
            )
        return run_configs
    else:
        ep, rep, epg, epc, eps = _build_engine_params(shared_args)
        return [
            RunConfig(
                run_idx=0,
                result_dir=shared_args.result_dir,
                run_label=shared_args.model_dir_name or "run-0",
                args=shared_args,
                engine_params=ep,
                reflection_engine_params=rep,
                engine_params_for_grounding=epg,
                engine_params_for_coding=epc,
                engine_params_for_searcher=eps,
            )
        ]


def run_env_tasks(
    task_queue: Queue,
    run_configs: list,
    shared_args: argparse.Namespace,
    per_run_scores: dict,
):
    active_environments = []
    env = None
    proc_name = current_process().name
    # Index run_configs by run_idx for fast lookup
    rc_by_idx = {rc.run_idx: rc for rc in run_configs}
    try:
        # Use IMAGE_ID_MAP for AWS provider to get snapshot_name
        snapshot_name = None
        region = getattr(shared_args, "region", None)
        if shared_args.provider_name == "aws" and region is not None:
            try:
                from desktop_env.providers.aws.manager import IMAGE_ID_MAP

                screen_size = (shared_args.screen_width, shared_args.screen_height)
                snapshot_name = IMAGE_ID_MAP[region].get(
                    screen_size, IMAGE_ID_MAP[region][(1920, 1080)]
                )
            except Exception as e:
                logger.error(f"Failed to get snapshot_name from IMAGE_ID_MAP: {e}")
                snapshot_name = None
        from mm_agents.vlaa_gui.agents.agent import Agent
        from mm_agents.vlaa_gui.agents.grounding import OSWorldACI
        from mm_agents.vlaa_gui.agents.verifier_agent import VerifierAgent

        # The first run config determines require_a11y_tree for the VM.
        # All configs that need a11y_tree will work because the VM enables
        # it at creation time; configs that don't need it simply ignore it.
        any_needs_a11y = any(
            rc.args.observation_type in ["a11y_tree", "screenshot_a11y_tree", "som"]
            for rc in run_configs
        )

        env = DesktopEnv(
            path_to_vm=shared_args.path_to_vm,
            action_space="pyautogui",
            provider_name=shared_args.provider_name,
            region=region,
            snapshot_name=snapshot_name,
            screen_size=(shared_args.screen_width, shared_args.screen_height),
            headless=shared_args.headless,
            os_type="Ubuntu",
            require_a11y_tree=any_needs_a11y,
            enable_proxy=True,
            client_password=getattr(shared_args, "client_password", ""),
        )

        agent_cache: dict[int, tuple] = {}

        def _get_or_create_agents(run_idx):
            if run_idx in agent_cache:
                return agent_cache[run_idx]
            rc = rc_by_idx[run_idx]
            a = rc.args

            coding_agent_flag = a.action_space == "pyautogui_coding"

            grounding_agent = OSWorldACI(
                env=env,
                platform="linux",
                engine_params_for_generation=rc.engine_params,
                engine_params_for_grounding=rc.engine_params_for_grounding,
                engine_params_for_searcher=rc.engine_params_for_searcher,
                width=shared_args.screen_width,
                height=shared_args.screen_height,
                resize_width=a.resize_width,
                grounding_model_type=a.grounding_model_type,
                code_agent_engine_params=rc.engine_params_for_coding,
                code_agent_budget=20,
            )
            agent = Agent(
                rc.engine_params,
                grounding_agent,
                platform="linux",
                action_space="pyautogui",
                observation_type=a.observation_type,
                with_reflection=a.with_reflection,
                search_engine=a.search_engine,
                memory_root_path="./agent_memory",
                memory_folder_name=a.kb_name,
                embedding_engine_type=a.embedding_engine_type,
                reflection_engine_params=rc.reflection_engine_params,
                memory_type=a.memory_type,
                coding_agent_flag=coding_agent_flag,
                max_image_in_trajectory=a.max_image_in_trajectory,
                debug=a.debug,
                enable_gate=a.enable_gate,
                loop_detection=a.loop_detection,
                feasibility_check=a.feasibility_check,
            )

            verifier = None
            if a.use_verifier:
                verifier = VerifierAgent(rc.engine_params, platform="linux")

            agent_cache[run_idx] = (agent, verifier, None)
            logger.info(
                f"[{proc_name}] Created agents for run {run_idx} ({rc.run_label})"
            )
            return agent_cache[run_idx]

        active_environments.append(env)
        proc_name = current_process().name
        
        logger.info(f"Process {proc_name} started.")
        while True:
            try:
                item = task_queue.get(timeout=5)
            except Exception:
                break
            run_idx, domain, example_id = item
            rc = rc_by_idx[run_idx]
            run_args = rc.args
            agent, verifier_agent, _ = _get_or_create_agents(run_idx)
            try:
                config_file = os.path.join(
                    shared_args.test_config_base_dir,
                    f"examples/{domain}/{example_id}.json",
                )
                with open(config_file, "r", encoding="utf-8") as f:
                    example = json.load(f)
                instruction = example["instruction"]
                
                example_result_dir = os.path.join(
                    rc.result_dir,
                    run_args.action_space,
                    run_args.observation_type,
                    run_args.model_dir_name,
                    domain,
                    example_id,
                )
                os.makedirs(example_result_dir, exist_ok=True)
                logger.info(f"[{proc_name}][Run {run_idx}][Domain]: {domain}")
                logger.info(f"[{proc_name}][Run {run_idx}][Example ID]: {example_id}")
                logger.info(f"[{proc_name}][Run {run_idx}][Instruction]: {instruction}")
                try:
                    lib_run_single.run_single_example_vlaa_gui(
                        agent,
                        env,
                        example,
                        run_args.max_steps,
                        instruction,
                        run_args,
                        example_result_dir,
                        per_run_scores[run_idx],
                        verifier_agent=verifier_agent,
                    )
                except Exception as e:
                    import traceback

                    logger.error(
                        f"Exception in {proc_name} run {run_idx} {domain}/{example_id}: {e}"
                    )
                    logger.error(traceback.format_exc())
                    try:
                        env.controller.end_recording(
                            os.path.join(example_result_dir, "recording.mp4")
                        )
                    except Exception as rec_e:
                        logger.error(f"Failed to end recording: {rec_e}")
                    with open(os.path.join(example_result_dir, "traj.jsonl"), "a") as f:
                        f.write(json.dumps({"Error": f"{domain}/{example_id} - {e}"}))
                        f.write("\n")
                
            except Exception as e:
                logger.error(f"Task-level error in {proc_name}: {e}")
                import traceback

                logger.error(traceback.format_exc())
    except Exception as e:
        logger.error(f"Process-level error in {proc_name}: {e}")
        import traceback

        logger.error(traceback.format_exc())
    finally:
        logger.info(f"{proc_name} cleaning up environment...")
        try:
            if env:
                env.close()
                logger.info(f"{proc_name} environment closed successfully")
        except Exception as e:
            logger.error(f"{proc_name} error during environment cleanup: {e}")


def config() -> argparse.Namespace:
    REASONING_CHOICES = ["none", "low", "medium", "high", "xhigh"]
    parser = argparse.ArgumentParser(
        description="Run end-to-end evaluation on OSWorld with the VLAA agent"
    )

    parser.add_argument(
        "--result_dir",
        type=str,
        default="results",
    )
    parser.add_argument(
        "--num_envs",
        type=int,
        default=30,
        help="Number of environments to run in parallel.",
    )
    parser.add_argument(
        "--test_all_meta_path",
        type=str,
        default="evaluation_examples/test_nogdrive.json",
    )
    parser.add_argument("--max_steps", type=int, default=100)
    parser.add_argument(
        "--headless",
        action="store_true",
        help="Run the environment headlessly.",
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Enable debug mode with verbose logging.",
    )
    parser.add_argument(
        "--provider_name",
        type=str,
        default="aws",
        help="Virtualization provider (vmware, docker, aws, azure, gcp, virtualbox).",
    )
    parser.add_argument(
        "--region", type=str, default="us-east-1", help="AWS region for the VM"
    )
    parser.add_argument(
        "--client_password", type=str, default="", help="Client password"
    )
    parser.add_argument("--path_to_vm", type=str, default=None)
    parser.add_argument("--screen_width", type=int, default=1920)
    parser.add_argument("--screen_height", type=int, default=1080)
    parser.add_argument("--sleep_after_execution", type=float, default=1.0)
    parser.add_argument("--domain", type=str, default="all")
    parser.add_argument(
        "--test_config_base_dir", type=str, default="evaluation_examples"
    )
    parser.add_argument("--max_trajectory_length", type=int, default=8)

    parser.add_argument("--model_provider", type=str, default="openai")
    parser.add_argument("--model", type=str, default="")
    parser.add_argument(
        "--model_dir_name",
        type=str,
        default="",
        help="Optional result directory name override for the model.",
    )
    parser.add_argument("--model_url", type=str, default="")
    parser.add_argument("--model_api_key", type=str, default="")
    parser.add_argument(
        "--model_api_keys",
        type=str,
        default="",
        help="Comma-separated fallback API keys.",
    )
    parser.add_argument("--model_temperature", type=float, default=1.0)
    parser.add_argument("--max_tokens", type=int, default=None)
    parser.add_argument("--model_project_id", type=str, default="")
    parser.add_argument("--model_region", type=str, default="")
    parser.add_argument("--api_version", type=str, default="")
    parser.add_argument("--model_thinking", action="store_true")
    parser.add_argument(
        "--model_thinking_level",
        type=str,
        choices=REASONING_CHOICES,
        default=None,
    )
    parser.add_argument("--model_thinking_budget", type=int, default=None)
    parser.add_argument(
        "--model_reasoning_effort",
        type=str,
        choices=REASONING_CHOICES,
        default=None,
    )
    parser.add_argument("--model_include_thoughts", action="store_true")

    parser.add_argument("--grounding_model_provider", type=str, default="openai")
    parser.add_argument("--grounding_model", type=str, default="")
    parser.add_argument("--grounding_model_url", type=str, default="")
    parser.add_argument("--grounding_model_api_key", type=str, default="")
    parser.add_argument("--grounding_width", type=int, default=None)
    parser.add_argument("--grounding_height", type=int, default=None)
    parser.add_argument("--resize_width", type=int, default=None)
    parser.add_argument("--grounding_model_region", type=str, default="")
    parser.add_argument("--enable_zoom_grounding", action="store_true")
    parser.add_argument("--zoom_grounding_crop_ratio", type=float, default=0.5)
    parser.add_argument("--grounding_model_type", type=str, default="single")
    parser.add_argument("--grounding_thinking", action="store_true")
    parser.add_argument("--grounding_temperature", type=float, default=None)
    parser.add_argument("--grounding_top_p", type=float, default=None)
    parser.add_argument("--grounding_thinking_budget", type=int, default=None)
    parser.add_argument(
        "--grounding_thinking_level",
        type=str,
        choices=REASONING_CHOICES,
        default=None,
    )
    parser.add_argument(
        "--grounding_reasoning_effort",
        type=str,
        choices=REASONING_CHOICES,
        default=None,
    )
    parser.add_argument("--grounding_include_thoughts", action="store_true")

    parser.add_argument("--endpoint_provider", type=str, default="")
    parser.add_argument("--endpoint_url", type=str, default="")
    parser.add_argument("--endpoint_api_key", type=str, default="")

    parser.add_argument("--coding_model_provider", type=str, default="openai")
    parser.add_argument("--coding_model", type=str, default="")
    parser.add_argument("--coding_model_url", type=str, default="")
    parser.add_argument("--coding_model_api_key", type=str, default="")
    parser.add_argument(
        "--coding_model_api_keys",
        type=str,
        default="",
        help="Comma-separated fallback API keys for the coding model.",
    )
    parser.add_argument("--coding_model_temperature", type=float, default=None)
    parser.add_argument("--coding_model_top_p", type=float, default=None)
    parser.add_argument("--coding_model_thinking", action="store_true")
    parser.add_argument("--coding_model_thinking_budget", type=int, default=None)
    parser.add_argument(
        "--coding_model_thinking_level",
        type=str,
        choices=REASONING_CHOICES,
        default=None,
    )
    parser.add_argument("--coding_model_include_thoughts", action="store_true")

    parser.add_argument("--embedding_engine_type", type=str, default="openai")
    parser.add_argument(
        "--observation_type",
        type=str,
        default="screenshot",
    )
    parser.add_argument("--planner_mode", type=str, default="proactive")
    parser.add_argument("--planner_hierarchical_depth", type=int, default=1)
    parser.add_argument("--with_reflection", action="store_true")
    parser.add_argument(
        "--reflection_thinking",
        action=argparse.BooleanOptionalAction,
        default=None,
    )
    parser.add_argument("--search_engine", type=str, default=None)
    parser.add_argument("--kb_name", type=str, default="agent_memory")
    parser.add_argument("--memory_type", type=str, default="null")
    parser.add_argument("--lexical_weight", type=float, default=0.0)

    parser.add_argument("--searcher_type", type=str, default="llm")
    parser.add_argument("--searcher_provider", type=str, default="gemini")
    parser.add_argument("--searcher_model", type=str, default="gemini-3.1-pro-preview")
    parser.add_argument("--searcher_api_key", type=str, default="")
    parser.add_argument("--searcher_url", type=str, default="")
    parser.add_argument("--searcher_budget", type=int, default=20)
    parser.add_argument("--searcher_temperature", type=float, default=None)
    parser.add_argument("--searcher_top_p", type=float, default=None)
    parser.add_argument(
        "--searcher_reasoning_effort",
        type=str,
        choices=REASONING_CHOICES,
        default=None,
    )

    parser.add_argument("--enable_gate", action="store_true")
    parser.add_argument("--loop_detection", action="store_true")
    parser.add_argument("--feasibility_check", action="store_true")
    parser.add_argument("--action_space", type=str, default="pyautogui_coding")
    parser.add_argument("--max_image_in_trajectory", type=int, default=8)
    parser.add_argument("--use_verifier", action="store_true")
    parser.add_argument("--enable_click_validation", action="store_true")

    parser.add_argument(
        "--runs-file",
        type=str,
        default=None,
        help=(
            "Path to a JSON file containing a list of per-run CLI override "
            "objects. Shared CLI args are used as defaults and each object "
            "overrides fields for one run."
        ),
    )

    return parser.parse_args()


def test(run_configs: list, shared_args: argparse.Namespace, all_tasks: list) -> None:
    global processes
    for rc in run_configs:
        logger.info("Run %d args: %s", rc.run_idx, rc.args)
    logger.info(f"Total tasks across all runs: {len(all_tasks)}")

    with Manager() as manager:
        # Per-run shared data structures
        per_run_scores: dict[int, Any] = {
            rc.run_idx: manager.list() for rc in run_configs
        }
        task_queue = manager.Queue()
        for item in all_tasks:
            task_queue.put(item)

        num_envs = shared_args.num_envs
        processes = []
        # Stagger process launches to avoid AWS RequestLimitExceeded errors.
        launch_delay = 2 if num_envs > 10 else 0
        worker_args = (
            task_queue,
            run_configs,
            shared_args,
            per_run_scores,
        )
        for i in range(num_envs):
            p = Process(
                target=run_env_tasks,
                args=worker_args,
                name=f"EnvProcess-{i + 1}",
            )
            p.daemon = True
            p.start()
            processes.append(p)
            logger.info(f"Started process {p.name} with PID {p.pid}")
            if launch_delay and i < num_envs - 1:
                time.sleep(launch_delay)

        try:
            while True:
                alive_count = 0
                for idx, p in enumerate(processes):
                    if not p.is_alive():
                        logger.warning(f"Process {p.name} died, restarting...")
                        new_p = Process(
                            target=run_env_tasks,
                            args=worker_args,
                            name=f"EnvProcess-Restart-{idx + 1}",
                        )
                        new_p.daemon = True
                        new_p.start()
                        processes[idx] = new_p
                        logger.info(
                            f"Restarted process {new_p.name} with PID {new_p.pid}"
                        )
                    else:
                        alive_count += 1

                if task_queue.empty():
                    logger.info("All tasks finished.")
                    break
                if alive_count == 0:
                    logger.error("All processes died, exiting.")
                    break
                time.sleep(5)
            for p in processes:
                p.join()

        except KeyboardInterrupt:
            logger.info(
                "Main process received KeyboardInterrupt. Initiating graceful shutdown..."
            )
            raise
        except Exception as e:
            logger.error(
                f"Unexpected error while waiting for processes: {e}", exc_info=True
            )
            for p in processes:
                if p.is_alive():
                    try:
                        logger.info(f"Terminating process {p.name} due to error...")
                        p.terminate()
                    except Exception as term_e:
                        logger.error(f"Error terminating process {p.name}: {term_e}")
            raise


def get_unfinished(
    action_space, use_model, observation_type, result_dir, total_file_json
):
    target_dir = os.path.join(result_dir, action_space, observation_type, use_model)

    if not os.path.exists(target_dir):
        return total_file_json

    finished = {}
    for domain in os.listdir(target_dir):
        finished[domain] = []
        domain_path = os.path.join(target_dir, domain)
        if os.path.isdir(domain_path):
            for example_id in os.listdir(domain_path):
                if example_id == "onboard":
                    continue
                example_path = os.path.join(domain_path, example_id)
                if os.path.isdir(example_path):
                    if "result.txt" not in os.listdir(example_path):
                        # empty all files under example_id
                        for file in os.listdir(example_path):
                            os.remove(os.path.join(example_path, file))
                    else:
                        finished[domain].append(example_id)

    if not finished:
        return total_file_json

    for domain, examples in finished.items():
        if domain in total_file_json:
            total_file_json[domain] = [
                x for x in total_file_json[domain] if x not in examples
            ]

    return total_file_json


def get_result(action_space, use_model, observation_type, result_dir, total_file_json):
    target_dir = os.path.join(result_dir, action_space, observation_type, use_model)
    if not os.path.exists(target_dir):
        print("New experiment, no result yet.")
        return None

    all_result = []

    for domain in os.listdir(target_dir):
        domain_path = os.path.join(target_dir, domain)
        if os.path.isdir(domain_path):
            for example_id in os.listdir(domain_path):
                example_path = os.path.join(domain_path, example_id)
                if os.path.isdir(example_path):
                    if "result.txt" in os.listdir(example_path):
                        # empty all files under example_id
                        try:
                            all_result.append(
                                float(
                                    open(
                                        os.path.join(example_path, "result.txt"), "r"
                                    ).read()
                                )
                            )
                        except Exception as e:
                            print(
                                f"Error reading result for {domain}/{example_id}: {e}"
                            )
                            all_result.append(0.0)

    if not all_result:
        print("New experiment, no result yet.")
        return None
    else:
        print("Current Success Rate:", sum(all_result) / len(all_result) * 100, "%")
        return all_result


if __name__ == "__main__":
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    ####### The complete version of the list of examples #######
    os.environ["TOKENIZERS_PARALLELISM"] = "false"
    args = config()

    # Build RunConfig list (single entry for legacy CLI, N entries for --runs-file)
    run_configs = build_run_configs(args)
    logger.info(f"Number of run configs: {len(run_configs)}")
    for rc in run_configs:
        logger.info(
            f"  Run {rc.run_idx}: label={rc.run_label} result_dir={rc.result_dir}"
        )

    # Save per-run args.json
    for rc in run_configs:
        path_to_args = os.path.join(
            rc.result_dir,
            rc.args.action_space,
            rc.args.observation_type,
            rc.args.model_dir_name,
            "args.json",
        )
        os.makedirs(os.path.dirname(path_to_args), exist_ok=True)
        with open(path_to_args, "w", encoding="utf-8") as f:
            json.dump(vars(rc.args), f, indent=4)

    # Load task metadata
    with open(args.test_all_meta_path, "r", encoding="utf-8") as f:
        test_all_meta = json.load(f)

    if args.domain != "all":
        test_all_meta = {args.domain: test_all_meta[args.domain]}

    # Total task count per run (all tasks, including finished) for milestones
    tasks_per_run = sum(len(v) for v in test_all_meta.values())
    total_task_count = tasks_per_run * len(run_configs)

    # Build per-run unfinished task lists and interleave them round-robin
    per_run_tasks: dict[int, list] = {}
    for rc in run_configs:
        # Deep-copy so get_unfinished can mutate freely
        meta_copy = {d: list(ids) for d, ids in test_all_meta.items()}
        unfinished = get_unfinished(
            rc.args.action_space,
            rc.args.model_dir_name,
            rc.args.observation_type,
            rc.result_dir,
            meta_copy,
        )
        per_run_tasks[rc.run_idx] = distribute_tasks(unfinished)

        left_info = ""
        for domain in unfinished:
            left_info += f"{domain}: {len(unfinished[domain])}\n"
        logger.info(f"[Run {rc.run_idx} ({rc.run_label})] Left tasks:\n{left_info}")

        get_result(
            rc.args.action_space,
            rc.args.model_dir_name,
            rc.args.observation_type,
            rc.result_dir,
            meta_copy,
        )

    # Interleave tasks round-robin across runs for even progress
    all_tasks: list = []
    max_task_len = max((len(tasks) for tasks in per_run_tasks.values()), default=0)
    for i in range(max_task_len):
        for run_idx in sorted(per_run_tasks.keys()):
            tasks = per_run_tasks[run_idx]
            if i < len(tasks):
                domain, example_id = tasks[i]
                all_tasks.append((run_idx, domain, example_id))

    logger.info(f"Total interleaved tasks to run: {len(all_tasks)}")
    test(run_configs, args, all_tasks)
