"""
Parallel OSWorld benchmark with on-the-fly env recovery.

Each worker process owns a persistent DesktopEnv (VM) and consumes tasks
from a shared queue.  Before every task the env is validated (snapshot
revert, Flask readiness, RDDS health + screenshot); if any check fails
the VM is torn down and recreated transparently.

Per-task progress is written to status.json *before* the agent starts, so
a trajectory ID is always recoverable even if the process crashes.

Usage:
    python mm_agents/surferH/run_benchmark_with_retries.py --num_envs 5
"""
import argparse
import datetime
import json
import logging
import os
import signal
import sys
import time
import traceback
from multiprocessing import Manager, Process, Queue, current_process
from pathlib import Path

import boto3
import requests
from rich.console import Console
from rich.progress import (
    BarColumn,
    MofNCompleteColumn,
    Progress,
    SpinnerColumn,
    TextColumn,
    TimeElapsedColumn,
    TimeRemainingColumn,
)

console = Console()

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from desktop_env.desktop_env import DesktopEnv
from lib_results_logger import log_task_completion
from mm_agents.surferH.install_rdds import install_rdds, wait_for_flask
from mm_agents.surferH.surfer_agent import SurferAgent

processes: list[Process] = []
is_terminating = False

RDDS_PORT = 8087
ACTION_SPACE = "pyautogui"
OBSERVATION_TYPE = "screenshot"

MAX_ENV_ATTEMPTS = 2
MAX_FRESHEN_ATTEMPTS = 3
MAX_POLL_ATTEMPTS = 3
POLL_BACKOFF_S = 15

SKIP_TASK_IDS: set[str] = {
    "78aed49a-a710-4321-a793-b611a7c5b56b",
    "897e3b53-5d4d-444b-85cb-2cdc8a97d903",
    "46407397-a7d5-4c6b-92c6-dbe038b1457b",
    "a0b9dc9c-fc07-4a88-8c5d-5e3ecad91bcb",
    "b52b40a5-ad70-4c53-b5b0-5650a8387052",
    "4e9f0faf-2ecc-4ae8-a804-28c9a75d1ddc",
    "0c825995-5b70-4526-b663-113f4c999dd2",
    "22a4636f-8179-4357-8e87-d1743ece1f81",
}

logger = logging.getLogger("surferH.benchmark_retry")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run OSWorld benchmark")
    parser.add_argument("--num_envs", type=int, default=1)
    parser.add_argument("--test_config_base_dir", type=str, default="evaluation_examples")
    parser.add_argument("--test_all_meta_path", type=str, default="evaluation_examples/test_all.json")
    parser.add_argument("--result_dir", type=str, default="./results")
    parser.add_argument("--region", type=str, default="us-west-2")
    parser.add_argument("--provider_name", type=str, default="aws")
    parser.add_argument("--headless", action="store_true", default=True)
    parser.add_argument("--screen_width", type=int, default=1920)
    parser.add_argument("--screen_height", type=int, default=1080)
    parser.add_argument("--log_level", type=str, default="INFO")
    parser.add_argument("--agent_identifier", type=str, default="OSWorldSurfer-Holo3-35B-A3B-20260330")
    return parser.parse_args()


def _apply_ec2_tags(instance_id: str, region: str) -> None:
    boto3.client("ec2", region_name=region).create_tags(
        Resources=[instance_id],
        Tags=[
            {"Key": "Project", "Value": "AgentEvaluation"},
            {"Key": "Name", "Value": "OSWorld"},
        ],
    )


# ---------------------------------------------------------------------------
# TaskWorker -- one instance per parallel process
# ---------------------------------------------------------------------------

class TaskWorker:
    """Owns a DesktopEnv reused across tasks, with on-the-fly recovery."""

    def __init__(self, args: argparse.Namespace, worker_id: str):
        self.args = args
        self.worker_id = worker_id
        self.env: DesktopEnv | None = None
        self.agent = SurferAgent(agent_identifier=args.agent_identifier)
        self.env_kwargs = self._build_env_kwargs()
        self.server_url: str | None = None
        self._create_env()

    # -- env lifecycle ------------------------------------------------------

    def _build_env_kwargs(self) -> dict:
        screen_size = (self.args.screen_width, self.args.screen_height)
        env_kwargs: dict = {
            "action_space": ACTION_SPACE,
            "provider_name": self.args.provider_name,
            "screen_size": screen_size,
            "headless": self.args.headless,
            "os_type": "Ubuntu",
        }
        if self.args.provider_name == "aws":
            from desktop_env.providers.aws.manager import IMAGE_ID_MAP

            ami_id = IMAGE_ID_MAP[self.args.region][screen_size]
            env_kwargs["region"] = self.args.region
            env_kwargs["snapshot_name"] = ami_id
        return env_kwargs

    def _create_env(self) -> None:
        if self.env:
            try:
                self.env.close()
            except Exception:
                pass
        self.env = DesktopEnv(**self.env_kwargs)
        _apply_ec2_tags(self.env.path_to_vm, self.args.region)
        logger.info(f"[{self.worker_id}] VM up: {self.env.vm_ip}")
        self.server_url = None

    def _recreate_env(self) -> None:
        logger.warning(f"[{self.worker_id}] Recreating environment...")
        self._create_env()

    def close(self) -> None:
        if self.env:
            try:
                self.env.close()
                logger.info(f"[{self.worker_id}] VM closed.")
            except Exception:
                pass

    # -- env readiness checks -----------------------------------------------

    def _check_rdds(self) -> None:
        private_url = f"http://{self.env.vm_ip}:{RDDS_PORT}"
        resp = requests.get(f"{private_url}/health", timeout=10)
        if resp.status_code != 200:
            raise RuntimeError(f"RDDS /health check failed: {resp.status_code}")
        resp = requests.get(f"{private_url}/screenshot", timeout=10)
        if resp.status_code != 200:
            raise RuntimeError(f"RDDS /screenshot check failed: {resp.status_code}")

    def _poll_services(self, task_id: str) -> str:
        """Poll for Flask readiness, install RDDS, poll for RDDS readiness.

        Retries each polling step up to MAX_POLL_ATTEMPTS before giving up.
        Returns the RDDS public server URL.
        """
        tag = f"[{self.worker_id}][{task_id[:8]}]"
        last_error: Exception | None = None

        for poll in range(MAX_POLL_ATTEMPTS):
            try:
                wait_for_flask(self.env)
                break
            except Exception as e:
                last_error = e
                logger.warning(f"{tag} Flask not ready (poll {poll + 1}/{MAX_POLL_ATTEMPTS})")
                if poll < MAX_POLL_ATTEMPTS - 1:
                    time.sleep(POLL_BACKOFF_S)
        else:
            raise last_error

        server_url = install_rdds(self.env, RDDS_PORT)
        logger.info(f"{tag} RDDS at {server_url}")

        for poll in range(MAX_POLL_ATTEMPTS):
            try:
                self._check_rdds()
                logger.info(f"{tag} RDDS health + screenshot OK")
                break
            except Exception as e:
                last_error = e
                logger.warning(f"{tag} RDDS not ready (poll {poll + 1}/{MAX_POLL_ATTEMPTS})")
                if poll < MAX_POLL_ATTEMPTS - 1:
                    time.sleep(POLL_BACKOFF_S)
        else:
            raise last_error

        return server_url

    def _prepare_env(self, example: dict) -> str:
        """Validate env readiness and set up the task.

        Three nested retry levels:
        1. MAX_ENV_ATTEMPTS  -- recreate VM from scratch
        2. MAX_FRESHEN_ATTEMPTS -- snapshot revert + reinstall RDDS
        3. MAX_POLL_ATTEMPTS -- poll Flask / RDDS readiness

        Returns the RDDS public server URL.
        """
        task_id = example["id"]
        tag = f"[{self.worker_id}][{task_id[:8]}]"
        last_error: Exception | None = None

        for env_attempt in range(MAX_ENV_ATTEMPTS):
            if env_attempt > 0:
                self._recreate_env()

            for freshen in range(MAX_FRESHEN_ATTEMPTS):
                try:
                    logger.info(f"{tag} Reverting VM snapshot...")
                    self.env.reset()
                    _apply_ec2_tags(self.env.path_to_vm, self.args.region)

                    server_url = self._poll_services(task_id)

                    logger.info(f"{tag} Running task setup...")
                    self.env.reset(task_config=example)
                    _apply_ec2_tags(self.env.path_to_vm, self.args.region)

                    server_url = self._poll_services(task_id)

                    self.server_url = server_url
                    return server_url
                except Exception as e:
                    last_error = e
                    logger.error(
                        f"{tag} Env prep failed "
                        f"(env {env_attempt + 1}/{MAX_ENV_ATTEMPTS}, "
                        f"freshen {freshen + 1}/{MAX_FRESHEN_ATTEMPTS}): {e}"
                    )
                    logger.error(traceback.format_exc())

        raise last_error

    # -- task execution -----------------------------------------------------

    @staticmethod
    def _write_status(result_dir: str, data: dict) -> None:
        with open(os.path.join(result_dir, "status.json"), "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2)

    def _run_task(
        self,
        example: dict,
        example_result_dir: str,
        shared_scores: list,
    ) -> None:
        task_id = example["id"]
        start_time = time.time()

        status: dict = {
            "task_id": task_id,
            "instruction": example["instruction"],
            "status": "running",
            "started_at": datetime.datetime.now().isoformat(),
        }
        self._write_status(example_result_dir, status)

        self.agent.reset(rdds_url=self.server_url)

        def on_trajectory_created(trajectory_id: str) -> None:
            status["trajectory_id"] = trajectory_id
            self._write_status(example_result_dir, status)

        message, actions = self.agent.predict(
            example["instruction"], {}, on_trajectory_created=on_trajectory_created,
        )
        elapsed = time.time() - start_time
        logger.info(f"[{self.worker_id}][{task_id[:8]}] AGP done in {elapsed:.0f}s: {message[:120]}")

        if "FAIL" in actions:
            self.env.action_history.append("FAIL")
        else:
            self.env.action_history.append("DONE")

        time.sleep(5)
        score = self.env.evaluate()
        logger.info(f"[{self.worker_id}][{task_id[:8]}] Score: {score}")

        status.update({
            "status": "completed",
            "score": score,
            "elapsed_s": elapsed,
            "agp_message": message,
            "agp_actions": actions,
        })
        self._write_status(example_result_dir, status)

        with open(os.path.join(example_result_dir, "result.txt"), "w", encoding="utf-8") as f:
            f.write(f"{score}\n")

        log_task_completion(example, score, example_result_dir, self.args)
        shared_scores.append(score)

    # -- main loop ----------------------------------------------------------

    def run(
        self,
        task_queue: Queue,
        shared_scores: list,
        shared_errors: list,
    ) -> None:
        while True:
            try:
                item = task_queue.get(timeout=5)
            except Exception:
                break

            domain, example_id = item
            config_file = os.path.join(
                self.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"[{self.worker_id}] {domain}/{example_id}: {example['instruction'][:80]}")

            example_result_dir = os.path.join(
                self.args.result_dir, ACTION_SPACE, OBSERVATION_TYPE,
                self.agent.agent_identifier, domain, example_id,
            )
            os.makedirs(example_result_dir, exist_ok=True)

            self.agent.last_trajectory_id = None
            try:
                self._prepare_env(example)
                self._run_task(example, example_result_dir, shared_scores)
            except Exception as e:
                logger.error(f"[{self.worker_id}] {domain}/{example_id} failed: {e}")
                logger.error(traceback.format_exc())
                shared_errors.append(f"{domain}/{example_id}: {e}")

                error_status: dict = {
                    "task_id": example_id,
                    "instruction": example["instruction"],
                    "status": "error",
                    "error": str(e),
                    "timestamp": datetime.datetime.now().isoformat(),
                }
                if self.agent.last_trajectory_id:
                    error_status["trajectory_id"] = self.agent.last_trajectory_id
                self._write_status(example_result_dir, error_status)


# ---------------------------------------------------------------------------
# Worker process entry point
# ---------------------------------------------------------------------------

def _worker_entry(
    task_queue: Queue,
    args: argparse.Namespace,
    shared_scores: list,
    shared_errors: list,
) -> None:
    logging.basicConfig(
        level=getattr(logging, args.log_level.upper()),
        format="%(asctime)s %(name)s %(levelname)s %(message)s",
        filename=args._log_file,
        force=True,
    )
    sys.stdout = open(os.devnull, "w")

    worker = TaskWorker(args, current_process().name)
    try:
        worker.run(task_queue, shared_scores, shared_errors)
    except Exception as e:
        logger.error(f"[{current_process().name}] Process-level error: {e}")
        logger.error(traceback.format_exc())
    finally:
        worker.close()


# ---------------------------------------------------------------------------
# Signal handling
# ---------------------------------------------------------------------------

def signal_handler(signum: int, frame) -> None:
    global is_terminating, processes
    if is_terminating:
        return
    is_terminating = True
    logger.info(f"Received signal {signum}. Shutting down...")

    for p in processes:
        if p.is_alive():
            p.terminate()

    time.sleep(1)
    for p in processes:
        if p.is_alive():
            os.kill(p.pid, signal.SIGKILL)

    logger.info("Shutdown complete.")
    sys.exit(0)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> None:
    args = parse_args()

    args.result_dir = f"{args.result_dir}_{datetime.datetime.now():%Y%m%d_%H%M%S}"

    log_file = Path(args.result_dir) / "benchmark.log"
    log_file.parent.mkdir(parents=True, exist_ok=True)
    args._log_file = str(log_file)
    logging.basicConfig(
        level=getattr(logging, args.log_level.upper()),
        format="%(asctime)s %(name)s %(levelname)s %(message)s",
        filename=args._log_file,
        force=True,
    )

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

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

    all_tasks = [(domain, eid) for domain, eids in test_all_meta.items() for eid in eids]
    all_tasks = [(d, eid) for d, eid in all_tasks if eid not in SKIP_TASK_IDS]
    total_tasks = len(all_tasks)

    console.print(f"[bold]Total tasks: {total_tasks}[/bold]")
    for domain, eids in test_all_meta.items():
        count = sum(1 for eid in eids if eid not in SKIP_TASK_IDS)
        if count:
            console.print(f"  {domain}: [cyan]{count}[/cyan]")
    console.print(f"[bold]Result dir:[/bold] [dim]{args.result_dir}[/dim]")
    console.print(f"Logs [dim]→ {log_file}[/dim]\n")

    if not all_tasks:
        console.print("[yellow]No tasks to run.[/yellow]")
        return

    global processes
    processes = []

    with Manager() as manager:
        shared_scores = manager.list()
        shared_errors = manager.list()
        task_queue: Queue = manager.Queue()

        for item in all_tasks:
            task_queue.put(item)

        for i in range(args.num_envs):
            p = Process(
                target=_worker_entry,
                args=(task_queue, args, shared_scores, shared_errors),
                name=f"VM-{i + 1}",
                daemon=True,
            )
            p.start()
            processes.append(p)

        with Progress(
            SpinnerColumn(),
            TextColumn("[bold blue]Running[/bold blue]"),
            BarColumn(bar_width=40),
            MofNCompleteColumn(),
            TextColumn("•"),
            TimeElapsedColumn(),
            TextColumn("•"),
            TimeRemainingColumn(),
            TextColumn("•"),
            TextColumn(
                "[green]ok={task.fields[ok]}[/green]  "
                "[red]err={task.fields[err]}[/red]  "
                "ok_rate={task.fields[ok_rate]}"
            ),
            console=console,
        ) as progress:
            tid = progress.add_task("Running", total=total_tasks, ok=0, err=0, ok_rate="0.0%")

            while True:
                scored = len(shared_scores)
                errs = len(shared_errors)
                done = scored + errs
                successes = sum(1 for s in shared_scores if float(s) > 0)
                ok_rate = f"{successes / done * 100:.1f}%" if done > 0 else "0.0%"
                progress.update(tid, completed=done, ok=successes, err=errs, ok_rate=ok_rate)

                if all(not p.is_alive() for p in processes) and task_queue.empty():
                    break

                if not task_queue.empty():
                    for idx, p in enumerate(processes):
                        if not p.is_alive():
                            logger.warning(f"{p.name} died, restarting...")
                            new_p = Process(
                                target=_worker_entry,
                                args=(task_queue, args, shared_scores, shared_errors),
                                name=f"VM-restart-{idx + 1}",
                                daemon=True,
                            )
                            new_p.start()
                            processes[idx] = new_p

                time.sleep(5)

            scored = len(shared_scores)
            errs = len(shared_errors)
            done = scored + errs
            successes = sum(1 for s in shared_scores if float(s) > 0)
            ok_rate = f"{successes / done * 100:.1f}%" if done > 0 else "0.0%"
            progress.update(tid, completed=done, ok=successes, err=errs, ok_rate=ok_rate)

        for p in processes:
            p.join(timeout=10)

        all_scores = list(shared_scores)
        all_errors = list(shared_errors)

    summary = {
        "total_tasks": total_tasks,
        "scored": len(all_scores),
        "passed": sum(1 for s in all_scores if s > 0),
        "errors": len(all_errors),
        "score": sum(all_scores) / len(all_scores) if all_scores else 0,
    }
    with open(os.path.join(args.result_dir, "final_summary.json"), "w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2)

    if all_scores:
        avg = sum(all_scores) / len(all_scores)
        successes = sum(1 for s in all_scores if s > 0)
        console.print(
            f"\n[bold]{'=' * 60}[/bold]"
            f"\n[bold]Final Results[/bold]"
            f"\n  Total tasks:  {total_tasks}"
            f"\n  Scored:       {len(all_scores)}"
            f"\n  [bold]Score:[/bold]       {avg:.1%}"
            f"\n  [green]Passed:      {successes}[/green]/{len(all_scores)}"
            f"\n  [red]Errors:      {len(all_errors)}[/red]"
            f"\n[bold]{'=' * 60}[/bold]"
        )
    else:
        console.print("\n[yellow]No scores collected.[/yellow]")

    if all_errors:
        console.print("\n[bold red]Errored tasks:[/bold red]")
        for err in all_errors:
            console.print(f"  [red]•[/red] {err}")

    console.print(f"\n[bold green]Agent identifier:[/bold green] {args.agent_identifier}")
    console.print(f"\n[bold green]Results directory:[/bold green] {args.result_dir}")


if __name__ == "__main__":
    main()
