"""Parallel multi-env runner for the M3 agent (Anthropic-Messages backend).

Example:
```
ANTHROPIC_BASE_URL=https://api.anthropic.com \\
ANTHROPIC_API_KEY=sk-... \\
ANTHROPIC_MODEL=<model-name> \\
M3_THINKING_MODE=adaptive \\
PYTHONPATH=. python scripts/python/run_multienv_m3.py \\
    --headless \\
    --observation_type screenshot \\
    --provider_name aws \\
    --num_envs 5 \\
    --max_steps 100 \\
    --result_dir ./results \\
    --test_all_meta_path evaluation_examples/test_all.json
```

Adaptive thinking
-----------------
The example above exports M3_THINKING_MODE=adaptive so the M3 agent
injects thinking={'type': 'adaptive'} into every request (the model
decides whether and how long to think). Leave it unset to disable
thinking. When thinking is on, temperature is forced to 1 and top_p is
dropped (required by the Anthropic API).

Endpoint credentials come from the standard Anthropic SDK env vars
(``ANTHROPIC_BASE_URL`` / ``ANTHROPIC_API_KEY`` / ``ANTHROPIC_MODEL``) or
from the equivalent ``--base_url`` / ``--api_key`` / ``--model`` CLI flags.

A local ``.env`` file is auto-loaded on startup (if ``python-dotenv`` is
installed) so credentials can also be set there.
"""

from __future__ import annotations
import argparse
import datetime
import json
import logging
import os
import re
import sys
import threading
import signal
import time
from typing import List
from multiprocessing import Process, Manager
from multiprocessing import current_process


# Auto-load ``./.env`` (if present and ``python-dotenv`` is installed) so
# ``ANTHROPIC_*`` and other configuration can live in a local file instead
# of always being exported in the shell. Runs BEFORE the M3Agent import
# because M3Agent reads env vars at construction time.
if os.path.exists(".env"):
    try:
        from dotenv import load_dotenv
        load_dotenv()
    except ImportError:
        pass


import lib_run_single  # noqa: E402
from desktop_env.desktop_env import DesktopEnv  # noqa: E402
from mm_agents.m3 import M3Agent  # noqa: E402


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


def config() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run end-to-end OSWorld evaluation with the M3 agent."
    )

    # endpoint config
    parser.add_argument(
        "--base_url", type=str, default=os.environ.get("ANTHROPIC_BASE_URL"),
        help="Anthropic-compatible endpoint URL. Falls back to "
             "the ANTHROPIC_BASE_URL environment variable.",
    )
    parser.add_argument(
        "--api_key", type=str, default=os.environ.get("ANTHROPIC_API_KEY"),
        help="API key for the endpoint. Tokens starting with 'sk-' are "
             "sent via the x-api-key header; everything else is sent as "
             "Authorization: Bearer. Falls back to the ANTHROPIC_API_KEY "
             "environment variable.",
    )
    parser.add_argument(
        "--extra_headers", type=str, default=None,
        help="JSON dict of extra HTTP headers to send on every request, e.g. "
             "'{\"X-Custom-Header\": \"value\"}'. Useful for enterprise-gateway "
             "routing or beta-flag headers. Falls back to ANTHROPIC_EXTRA_HEADERS "
             "(also a JSON dict). Either source is forwarded to the Anthropic "
             "SDK's default_headers.",
    )

    # environment config
    parser.add_argument("--path_to_vm", type=str, default=None)
    parser.add_argument("--headless", action="store_true", help="Run in headless machine")
    parser.add_argument("--action_space", type=str, default="pyautogui", help="Action type")
    parser.add_argument(
        "--observation_type",
        choices=["screenshot", "a11y_tree", "screenshot_a11y_tree", "som"],
        default="screenshot",
        help="Observation type",
    )
    parser.add_argument("--sleep_after_execution", type=float, default=3.0)
    parser.add_argument("--max_steps", type=int, default=100)

    # evaluation config
    parser.add_argument("--test_config_base_dir", type=str, default="evaluation_examples")

    # lm config
    parser.add_argument(
        "--model", type=str, default=None,
        help="Model name. If omitted, falls back to ANTHROPIC_MODEL from config."
    )
    parser.add_argument("--temperature", type=float, default=0.6)
    parser.add_argument("--top_p", type=float, default=None)
    parser.add_argument("--max_tokens", type=int, default=8192)

    # M3 agent config
    parser.add_argument(
        "--coordinate_type", type=str, default="relative",
        choices=["relative", "absolute"],
        help="M3 emits [0, 1000] normalised coordinates; 'relative' is the "
             "default. Use 'absolute' if you feed the model a differently-sized "
             "image and want the parser to scale by the processed dimensions."
    )
    parser.add_argument(
        "--max_trajectory_length", type=int, default=10,
        help="Image-history budget. Mapped onto M3 only_n_most_recent_images."
    )
    parser.add_argument(
        "--platform", type=str, default="ubuntu", choices=["ubuntu", "macos", "windows"],
    )

    # example config
    parser.add_argument("--domain", type=str, default="all")
    parser.add_argument(
        "--test_all_meta_path", type=str, default="evaluation_examples/test_all.json",
    )

    # logging
    parser.add_argument("--result_dir", type=str, default="./results")
    parser.add_argument(
        "--num_envs", type=int, default=1,
        help="Number of environments to run in parallel"
    )
    parser.add_argument(
        "--env_start_delay", type=float, default=1.0,
        help="Seconds to sleep between consecutive worker spawns. Some cloud "
             "providers throttle bursty VM allocations; a small stagger "
             "(default 1.0 s) avoids that. Set to 0 to spawn all workers "
             "immediately."
    )
    parser.add_argument(
        "--log_level", type=str,
        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
        default="INFO", help="Set the logging level",
    )

    # provider config
    parser.add_argument("--region", type=str, default="us-east-1", help="AWS region for the VM")
    parser.add_argument(
        "--provider_name", type=str, default="aws",
        choices=["aws", "virtualbox", "vmware", "docker", "azure"],
        help="Provider name",
    )
    parser.add_argument(
        "--client_password", type=str, default="",
        help="Client VM password (sudo). Default 'osworld-public-evaluation' for AWS, "
             "'password' for docker/vmware.",
    )
    parser.add_argument("--screen_width", type=int, default=1920, help="Screen width")
    parser.add_argument("--screen_height", type=int, default=1080, help="Screen height")
    parser.add_argument(
        "--enable_proxy", action="store_true",
        help="Allow tasks whose config sets 'proxy: true' to use the proxy "
             "setup. When omitted (default), the proxy setup is skipped even "
             "for tasks that request it.",
    )

    args = parser.parse_args()
    return args


args = config()  # Get command line arguments first

logger = logging.getLogger()
log_level = getattr(logging, args.log_level.upper())
logger.setLevel(log_level)

datetime_str: str = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
os.makedirs("logs", exist_ok=True)

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

file_handler.setLevel(logging.INFO)
debug_handler.setLevel(logging.DEBUG)
stdout_handler.setLevel(log_level)

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"
)
file_handler.setFormatter(formatter)
debug_handler.setFormatter(formatter)
stdout_handler.setFormatter(formatter)

stdout_handler.addFilter(logging.Filter("desktopenv"))

logger.addHandler(file_handler)
logger.addHandler(debug_handler)
logger.addHandler(stdout_handler)

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


def distribute_tasks(test_all_meta: dict) -> List[tuple]:
    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 _resolved_model(args) -> str:
    """Model name as the API expects it (with bracket suffix preserved)."""
    return args.model or os.environ.get("ANTHROPIC_MODEL", "m3")


def _resolved_extra_headers(args):
    """Parse the --extra_headers JSON arg into a dict, or return None so
    M3Agent falls back to its own ANTHROPIC_EXTRA_HEADERS env-var resolution.
    """
    raw = getattr(args, "extra_headers", None)
    if not raw:
        return None
    try:
        parsed = json.loads(raw)
    except json.JSONDecodeError as e:
        logger.warning(
            "--extra_headers is not valid JSON (%s); ignoring CLI value, "
            "M3Agent may still load from ANTHROPIC_EXTRA_HEADERS.", e
        )
        return None
    if not isinstance(parsed, dict):
        logger.warning(
            "--extra_headers must decode to a dict, got %s; ignoring.",
            type(parsed).__name__,
        )
        return None
    return {str(k): str(v) for k, v in parsed.items()}


_PATH_SANITIZE_RE = re.compile(r"[^A-Za-z0-9._\-]+")


def _sanitize_for_path(name: str) -> str:
    """Make a model name safe to use as a filesystem path segment.

    Some model names (e.g. ``m3-model[1m]``) contain ``[``/``]`` which Linux
    filesystems accept but shell globs expand. Replace anything outside
    ``[A-Za-z0-9._-]`` with ``_`` and collapse runs.

    ``m3-model[1m]`` → ``m3-model_1m_`` → ``m3-model_1m``
    """
    s = _PATH_SANITIZE_RE.sub("_", name).strip("_")
    return s or "model"


def _resolved_password(args) -> str:
    if args.client_password:
        return args.client_password
    if args.provider_name == "aws":
        return "osworld-public-evaluation"
    return "password"


def _is_transient_infra_error(exc) -> bool:
    """Heuristic: True if the exception is from external infra that may
    recover on retry (cloud-provider capacity / API throttling / network
    blips), not from agent logic or model output. Used to decide whether
    to put a failed task back on the queue."""
    s = repr(exc) + " " + str(exc)
    markers = (
        # Networking blips (cloud-agnostic)
        "ConnectTimeout",
        "ConnectionError",
        "ReadTimeoutError",
        "Read timed out",
        "Connection refused",
        "Connection reset",
        # API throttling (botocore convention, used by AWS and others)
        "Throttling",
        "ThrottlingException",
        # AWS EC2 capacity / quota / availability
        "InsufficientInstanceCapacity",
        "RequestLimitExceeded",
        "InstanceLimitExceeded",
        "ServiceUnavailable",
        # OSWorld runner-level wrapper message
        "Failed to allocate",
    )
    return any(m in s for m in markers)


def _watchdog_loop(processes_ref, worker_state, task_queue,
                   threshold_s: int, max_requeue: int, poll_interval_s: int = 60):
    """Detect workers that are stuck on an unresponsive VM client.

    A live worker that's making real progress writes a new
    ``api_logs/NNN_<ts>.json`` file every LLM call (including each
    in-step retry attempt). If the newest file's mtime hasn't moved in
    ``threshold_s`` seconds, the env-side reset/step is almost certainly
    blocked on a terminated VM and won't recover on its own.

    When a stuck worker is detected:
      1. Re-queue the task with retry_count+1 (if still under
         ``max_requeue``) so it doesn't get silently lost
      2. SIGTERM the worker; if still alive after 10s, SIGKILL
      3. Let the main supervisor loop spawn a fresh worker (with a fresh
         DesktopEnv) to pick up the next task

    Workers register their current task via ``worker_state[pid] =
    (api_logs_dir, start_ts, domain, example_id, retry_count)`` before
    starting ``run_single_example`` and pop on exit.
    """
    import signal as _signal
    while not is_terminating:
        time.sleep(poll_interval_s)
        if is_terminating:
            break
        now = time.time()
        for pid, info in list(worker_state.items()):
            try:
                logdir, start_ts, domain, example_id, retry_count = info
            except (TypeError, ValueError):
                continue
            try:
                files = [f for f in os.listdir(logdir) if f.endswith(".json")]
                if files:
                    latest_mtime = max(
                        os.path.getmtime(os.path.join(logdir, f)) for f in files
                    )
                else:
                    latest_mtime = start_ts
            except OSError:
                latest_mtime = start_ts
            elapsed = now - latest_mtime
            if elapsed <= threshold_s:
                continue
            logger.warning(
                f"[watchdog] PID {pid} stuck on {domain}/{example_id} — "
                f"no api_log update in {elapsed:.0f}s (threshold {threshold_s}s); "
                f"requeueing + terminating worker"
            )
            if retry_count < max_requeue:
                try:
                    task_queue.put((domain, example_id, retry_count + 1))
                    logger.info(
                        f"[watchdog] requeued {domain}/{example_id} "
                        f"(retry {retry_count + 1}/{max_requeue})"
                    )
                except Exception as e:
                    logger.error(f"[watchdog] failed to requeue {domain}/{example_id}: {e}")
            else:
                logger.error(
                    f"[watchdog] {domain}/{example_id} exhausted requeue "
                    f"budget ({max_requeue}); giving up"
                )
            # Find and kill the matching Process (terminate → SIGKILL fallback)
            for p in processes_ref:
                if p is not None and p.pid == pid and p.is_alive():
                    try:
                        p.terminate()
                        p.join(timeout=10)
                        if p.is_alive():
                            logger.warning(
                                f"[watchdog] PID {pid} ignored SIGTERM; sending SIGKILL"
                            )
                            os.kill(p.pid, _signal.SIGKILL)
                    except Exception as e:
                        logger.error(f"[watchdog] failed to kill PID {pid}: {e}")
                    break
            worker_state.pop(pid, None)


def run_env_tasks(task_queue, args: argparse.Namespace, shared_scores: list,
                  worker_state=None):
    MAX_REQUEUE = int(os.environ.get("M3_TASK_MAX_REQUEUE", "3"))
    REQUEUE_BACKOFF_S = float(os.environ.get("M3_TASK_REQUEUE_BACKOFF_S", "5"))
    active_environments = []
    env = None
    try:
        client_password = _resolved_password(args)
        screen_size = (args.screen_width, args.screen_height)

        env_kwargs = dict(
            path_to_vm=args.path_to_vm,
            action_space=args.action_space,
            provider_name=args.provider_name,
            screen_size=screen_size,
            headless=args.headless,
            os_type="Ubuntu",
            require_a11y_tree=args.observation_type
                in ["a11y_tree", "screenshot_a11y_tree", "som"],
            client_password=client_password,
            enable_proxy=args.enable_proxy,
        )
        if args.provider_name == "aws":
            from desktop_env.providers.aws.manager import IMAGE_ID_MAP
            ami_id = IMAGE_ID_MAP[args.region].get(
                screen_size, IMAGE_ID_MAP[args.region][(1920, 1080)]
            )
            env_kwargs["region"] = args.region
            env_kwargs["snapshot_name"] = ami_id

        env = DesktopEnv(**env_kwargs)
        active_environments.append(env)

        logger.info(f"Process {current_process().name} started.")
        while True:
            try:
                item = task_queue.get(timeout=5)
            except Exception:
                break
            # Backwards-compatible: items can be (domain, example_id) or
            # (domain, example_id, retry_count). The retry_count ticks up
            # each time a transient infra error sends the task back to the
            # queue (see "transient infra" branch below).
            if len(item) == 3:
                domain, example_id, retry_count = item
            else:
                domain, example_id = item
                retry_count = 0
            try:
                config_file = os.path.join(
                    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)
                logger.info(f"[{current_process().name}][Domain]: {domain}")
                logger.info(f"[{current_process().name}][Example ID]: {example_id}")
                logger.info(
                    f"[{current_process().name}][Instruction]: {example['instruction']}"
                )
                example_result_dir = os.path.join(
                    args.result_dir,
                    args.action_space,
                    args.observation_type,
                    _sanitize_for_path(_resolved_model(args)),
                    domain,
                    example_id,
                )
                os.makedirs(example_result_dir, exist_ok=True)

                agent = M3Agent(
                    platform=args.platform,
                    model=_resolved_model(args),
                    base_url=args.base_url,
                    api_key=args.api_key,
                    extra_headers=_resolved_extra_headers(args),
                    max_tokens=args.max_tokens,
                    top_p=args.top_p,
                    temperature=args.temperature,
                    action_space=args.action_space,
                    observation_type=args.observation_type,
                    coordinate_type=args.coordinate_type,
                    max_trajectory_length=args.max_trajectory_length,
                    client_password=client_password,
                )
                api_logs_dir = os.path.join(example_result_dir, "api_logs")
                if hasattr(agent, "set_api_log_dir"):
                    agent.set_api_log_dir(api_logs_dir)

                # Register with the main-process watchdog so it can detect a
                # stuck VM (no new api_log written in M3_STUCK_THRESHOLD_S
                # seconds) and recycle this worker. Pop on every exit path
                # below so a finished/errored task doesn't keep us flagged.
                if worker_state is not None:
                    worker_state[os.getpid()] = (
                        api_logs_dir, time.time(), domain, example_id, retry_count,
                    )
                try:
                    try:
                        lib_run_single.run_single_example(
                            agent,
                            env,
                            example,
                            args.max_steps,
                            example["instruction"],
                            args,
                            example_result_dir,
                            shared_scores,
                        )
                    except Exception as e:
                        import traceback
                        # Transient infra failures (cloud-provider capacity,
                        # API throttling, network blips) are not the agent's
                        # fault — the task hasn't really had a fair shot. Put
                        # it back on the queue (with retry counter to avoid
                        # infinite loops) and clear the half-written result
                        # dir so the next worker starts clean. Without this,
                        # the task is silently dropped and never gets a
                        # result.txt.
                        if _is_transient_infra_error(e) and retry_count < MAX_REQUEUE:
                            logger.warning(
                                f"[requeue] {domain}/{example_id} hit transient infra "
                                f"({type(e).__name__}: {str(e)[:120]}) — retry "
                                f"{retry_count + 1}/{MAX_REQUEUE}"
                            )
                            try:
                                env.controller.end_recording(
                                    os.path.join(example_result_dir, "recording.mp4")
                                )
                            except Exception:
                                pass
                            try:
                                import shutil
                                shutil.rmtree(example_result_dir, ignore_errors=True)
                            except Exception as cleanup_e:
                                logger.warning(
                                    f"[requeue] cleanup failed for "
                                    f"{example_result_dir}: {cleanup_e}"
                                )
                            try:
                                time.sleep(REQUEUE_BACKOFF_S)
                                task_queue.put(
                                    (domain, example_id, retry_count + 1)
                                )
                            except Exception as putback_e:
                                logger.error(
                                    f"[requeue] failed to put back "
                                    f"{domain}/{example_id}: {putback_e}"
                                )
                            # Don't break out of the worker loop — the same
                            # worker can pick another task while we wait for
                            # the cloud API to recover quota / rate-limit.
                            continue
                        logger.error(
                            f"Exception in {current_process().name} "
                            f"{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")
                finally:
                    # Unflag this worker so the watchdog stops tracking the
                    # finished/errored task. Runs on every exit path of the
                    # inner try/except, including the requeue ``continue``.
                    if worker_state is not None:
                        worker_state.pop(os.getpid(), None)
            except Exception as e:
                logger.error(f"Task-level error in {current_process().name}: {e}")
                import traceback
                logger.error(traceback.format_exc())

    except Exception as e:
        logger.error(f"Process-level error in {current_process().name}: {e}")
        import traceback
        logger.error(traceback.format_exc())
    finally:
        logger.info(f"{current_process().name} cleaning up environment...")
        try:
            if env:
                env.close()
                logger.info(
                    f"{current_process().name} environment closed successfully"
                )
        except Exception as e:
            logger.error(
                f"{current_process().name} error during environment cleanup: {e}"
            )


def signal_handler(signum, frame):
    """Handle SIGINT/SIGTERM to gracefully shutdown environments."""
    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)


def test(args: argparse.Namespace, test_all_meta: dict) -> None:
    global processes
    logger.info("Args: %s", args)
    all_tasks = distribute_tasks(test_all_meta)
    logger.info(f"Total tasks: {len(all_tasks)}")
    # Tunables for the stuck-worker watchdog (see _watchdog_loop).
    STUCK_THRESHOLD_S = int(os.environ.get("M3_STUCK_THRESHOLD_S", "900"))
    WATCHDOG_POLL_S = int(os.environ.get("M3_WATCHDOG_POLL_S", "60"))
    MAX_REQUEUE = int(os.environ.get("M3_TASK_MAX_REQUEUE", "3"))
    with Manager() as manager:
        shared_scores = manager.list()
        task_queue = manager.Queue()
        # PID → (api_logs_dir, start_ts, domain, example_id, retry_count).
        # Workers update on task start/end; watchdog reads to detect stuck VMs.
        worker_state = manager.dict()
        for item in all_tasks:
            task_queue.put(item)
        num_envs = args.num_envs
        env_start_delay = max(0.0, float(getattr(args, "env_start_delay", 0.0)))
        processes = []
        for i in range(num_envs):
            p = Process(
                target=run_env_tasks,
                args=(task_queue, args, shared_scores, worker_state),
                name=f"EnvProcess-{i+1}",
            )
            p.daemon = True
            p.start()
            processes.append(p)
            logger.info(f"Started process {p.name} with PID {p.pid}")
            # Stagger worker spawns so all of them don't hit the cloud
            # provider's VM-allocation API at the same instant. Some
            # providers throttle bursty allocation requests and return
            # out-of-stock or throttling errors when too many calls fire
            # at once. Spreading the calls over ~num_envs * env_start_delay
            # seconds lets the allocator process them sequentially.
            if env_start_delay > 0 and i < num_envs - 1:
                time.sleep(env_start_delay)
        # Start the stuck-worker watchdog. Runs as a daemon thread in the
        # main process; exits when ``is_terminating`` flips on SIGINT/SIGTERM
        # or when ``test()`` returns.
        watchdog_thread = threading.Thread(
            target=_watchdog_loop,
            args=(processes, worker_state, task_queue,
                  STUCK_THRESHOLD_S, MAX_REQUEUE, WATCHDOG_POLL_S),
            name="StuckWorkerWatchdog",
            daemon=True,
        )
        watchdog_thread.start()
        logger.info(
            f"Started watchdog (threshold={STUCK_THRESHOLD_S}s, "
            f"poll={WATCHDOG_POLL_S}s, max_requeue={MAX_REQUEUE})"
        )
        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=(task_queue, args, shared_scores, worker_state),
                            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
        scores = list(shared_scores)
    logger.info(
        f"Average score: {sum(scores) / len(scores) if scores else 0}"
    )


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):
                        for file in os.listdir(example_path):
                            try:
                                os.remove(os.path.join(example_path, file))
                            except IsADirectoryError:
                                pass
                    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):
                        try:
                            all_result.append(
                                float(
                                    open(
                                        os.path.join(example_path, "result.txt"),
                                        "r",
                                    ).read()
                                )
                            )
                        except Exception:
                            all_result.append(0.0)

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


if __name__ == "__main__":
    os.environ["TOKENIZERS_PARALLELISM"] = "false"

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    try:
        model_for_paths = _sanitize_for_path(_resolved_model(args))

        path_to_args = os.path.join(
            args.result_dir,
            args.action_space,
            args.observation_type,
            model_for_paths,
            "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(args), f, indent=4)

        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]}

        test_file_list = get_unfinished(
            args.action_space,
            model_for_paths,
            args.observation_type,
            args.result_dir,
            test_all_meta,
        )
        left_info = ""
        for domain in test_file_list:
            left_info += f"{domain}: {len(test_file_list[domain])}\n"
        logger.info(f"Left tasks:\n{left_info}")

        get_result(
            args.action_space,
            model_for_paths,
            args.observation_type,
            args.result_dir,
            test_all_meta,
        )
        test(args, test_file_list)
    except KeyboardInterrupt:
        logger.info("Main process received KeyboardInterrupt.")
    except Exception as e:
        logger.error(f"Unexpected error in main process: {e}", exc_info=True)
        signal_handler(signal.SIGTERM, None)
    finally:
        logger.info("Main process final cleanup...")
        for env in active_environments:
            if env is not None:
                try:
                    logger.info("Closing environment in final cleanup...")
                    env.close()
                    logger.info("Environment closed successfully in final cleanup")
                except Exception as e:
                    logger.error(f"Error during final environment cleanup: {e}")

        for p in processes:
            if p is not None and p.is_alive():
                try:
                    logger.info(f"Terminating process {p.name}...")
                    p.terminate()
                except Exception as e:
                    logger.error(f"Error terminating process: {e}")

        time.sleep(1)

        for p in processes:
            if p is not None and p.is_alive():
                try:
                    logger.info(f"Force killing process {p.name}...")
                    os.kill(p.pid, signal.SIGKILL)
                    logger.info(f"Process {p.name} force killed")
                except Exception as e:
                    logger.error(f"Error force killing process: {e}")
