# Copyright (c) Metis. All rights reserved.

"""Orchestrator ABC and DockerOrchestrator implementation.

The ``Orchestrator`` base class defines the contract for compute lifecycle
management.  Each backend (Docker, Nomad, KubeRay) implements this contract.
``ComputeHandle`` is an opaque token that the orchestrator returns from
``start()`` and accepts back in every subsequent call.

Usage (via AlgorithmContext)::

    orch = DockerOrchestrator()
    handle = orch.start(
        run_id="run-abc",
        image="my-agent:v1",
        store_url="http://host.docker.internal:9876",
        count=3,
        lifecycle=ContainerLifecycle.PERSISTENT,
    )
    orch.wait_ready(handle, timeout=120.0)
    # ... algorithm enqueues rollouts, N agents process them ...
    orch.stop(handle)
"""

from __future__ import annotations

import logging
import time as _time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


# ═══════════════════════════════════════════════════════════════════════
# Enums & data classes
# ═══════════════════════════════════════════════════════════════════════


class ComputeStatus(str, Enum):
    """Lifecycle state of orchestrated compute."""

    PENDING = "pending"
    RUNNING = "running"
    SUCCEEDED = "succeeded"
    FAILED = "failed"
    UNKNOWN = "unknown"


class ContainerLifecycle(str, Enum):
    """How agent containers are recycled during a run.

    PERSISTENT: containers stay alive for the entire run duration.
               Best for stateless agents -- avoids startup overhead.
    EPHEMERAL:  containers are replaced after each rollout batch.
               Guarantees clean state for stateful agents (e.g. CRM sim).
    """

    PERSISTENT = "persistent"
    EPHEMERAL = "ephemeral"


@dataclass
class ComputeHandle:
    """Opaque handle returned by start(), passed back to stop()/status().

    Each orchestrator subclass stores whatever it needs here.
    Docker stores container IDs; Nomad stores a job ID; KubeRay stores a
    RayJob name; etc.
    """

    run_id: str
    backend: str  # "docker", "nomad", "kuberay", ...
    identifiers: Dict[str, str] = field(default_factory=dict)
    # e.g. {"agent-0": "abc123", "agent-1": "def456"}  (Docker)
    # e.g. {"job_id": "mantis-run-42"}                  (Nomad)
    # e.g. {"rayjob": "mantis-verl-run-42"}             (KubeRay)
    metadata: Dict[str, Any] = field(default_factory=dict)


# ═══════════════════════════════════════════════════════════════════════
# Orchestrator ABC
# ═══════════════════════════════════════════════════════════════════════


class Orchestrator(ABC):
    """Base class for compute orchestrators.

    An orchestrator manages the lifecycle of compute resources for a
    platform run.  Subclasses handle the specifics of Docker, Nomad,
    KubeRay, etc.

    Lifecycle (called by AlgorithmContext)::

        handle = orch.start(...)        # __aenter__
        orch.wait_ready(handle)         # block until healthy
        # ... algorithm runs ...
        orch.status(handle)             # health checks during run
        orch.stop(handle)               # __aexit__
    """

    @property
    @abstractmethod
    def available(self) -> bool:
        """Can this orchestrator function in the current environment?"""
        ...

    @property
    @abstractmethod
    def backend_name(self) -> str:
        """Short identifier: 'docker', 'nomad', 'kuberay', etc."""
        ...

    @abstractmethod
    def start(
        self,
        *,
        run_id: str,
        image: str,
        store_url: str,
        count: int = 1,
        lifecycle: ContainerLifecycle = ContainerLifecycle.PERSISTENT,
        env: Optional[Dict[str, str]] = None,
        labels: Optional[Dict[str, str]] = None,
        runner_config: Optional[Dict[str, Any]] = None,
        resource_limits: Optional[Dict[str, Any]] = None,
        # --- reactive scaling (optional) ---
        store: Optional[Any] = None,
        n_runners_per_unit: int = 1,
        max_units: int = 20,
        # --- extended fields (used by KubeRay, ignored by Docker/Nomad) ---
        trainer_resources: Optional[Dict[str, Any]] = None,
        services: Optional[Dict[str, Any]] = None,
        algorithm_config: Optional[Dict[str, Any]] = None,
    ) -> ComputeHandle:
        """Provision compute for a run.

        For Docker:  starts ``count`` agent containers
        For Nomad:   submits a job spec with task group ``count``
        For KubeRay: creates a RayJob CR (head + workers)

        If ``store`` is provided, the orchestrator registers a listener
        for reactive scaling (start containers when rollouts are enqueued)
        and ephemeral recycling (replace containers after rollout completion
        when lifecycle is EPHEMERAL).

        Args:
            run_id: MantisRun ID.
            image: Container image to run.
            store_url: HTTP URL of the worker's LightningStoreServer.
            count: Number of agent instances to start immediately.
                Use 0 for fully reactive (containers start on demand).
            lifecycle: Whether containers persist or are replaced per batch.
            env: Additional environment variables.
            labels: Metadata labels for the containers.
            runner_config: Runner lifecycle config from agent manifest.
            resource_limits: Container resource limits from agent manifest.
            store: LightningStore instance. If provided, enables reactive
                scaling via store listener hooks.
            n_runners_per_unit: Rollouts a single compute unit handles
                concurrently (used for capacity calculation).
            max_units: Hard cap on number of compute units.
            trainer_resources: GPU/memory for trainer (KubeRay only).
            services: Sidecar services (KubeRay only).
            algorithm_config: Algorithm-specific config (KubeRay only).

        Returns:
            ComputeHandle with identifiers for the provisioned compute.

        Raises:
            RuntimeError: If compute cannot be provisioned.
        """
        ...

    @abstractmethod
    def stop(self, handle: ComputeHandle, timeout: int = 30) -> None:
        """Tear down all compute for a run.

        Must be idempotent -- safe to call multiple times or after
        compute has already exited.
        """
        ...

    @abstractmethod
    def status(self, handle: ComputeHandle) -> ComputeStatus:
        """Check aggregate health of the compute for a run.

        Returns RUNNING if all members are healthy, FAILED if any have
        crashed, PENDING if still starting up.
        """
        ...

    @abstractmethod
    def wait_ready(self, handle: ComputeHandle, timeout: float = 120.0) -> None:
        """Block until all provisioned compute is healthy.

        Docker: poll container.status until all report 'running'.
        Nomad:  poll deployment health.
        KubeRay: poll RayJob status.

        Raises:
            RuntimeError: If timeout exceeded before all members are ready.
        """
        ...

    @abstractmethod
    def scale(self, handle: ComputeHandle, count: int) -> None:
        """Adjust the number of running instances.

        Docker: start/stop delta containers, update handle.identifiers.
        Nomad:  update job count, scheduler reconciles.

        Args:
            handle: ComputeHandle to scale.
            count: New desired count.
        """
        ...

    @abstractmethod
    def replace_member(
        self,
        handle: ComputeHandle,
        member_key: str,
    ) -> Optional[str]:
        """Replace one member of the compute group (ephemeral mode).

        Docker: stop container, start fresh one, update handle.
        Nomad:  ``nomad alloc restart``, or no-op (scheduler handles).

        Args:
            handle: ComputeHandle containing the member.
            member_key: Key in handle.identifiers (e.g. "agent-0").

        Returns:
            New container/alloc ID, or None if replacement failed.
        """
        ...

    def members(self, handle: ComputeHandle) -> List[Dict[str, Any]]:
        """List individual members of the compute group.

        Optional -- returns empty list by default.
        """
        return []

    def get_logs(self, handle: ComputeHandle, tail: int = 100) -> str:
        """Fetch recent logs. Optional -- returns empty string by default."""
        return ""


# ═══════════════════════════════════════════════════════════════════════
# DockerOrchestrator
# ═══════════════════════════════════════════════════════════════════════


class DockerOrchestrator(Orchestrator):
    """Manages agent containers via Docker SDK.

    For local development, uses the Docker socket to start/stop containers.
    Gracefully degrades if Docker is not available (e.g., in CI without Docker).
    """

    def __init__(self, docker_host: Optional[str] = None):
        """
        Args:
            docker_host: Docker daemon URL. Defaults to unix:///var/run/docker.sock.
        """
        self._client = None
        self._stop_timeouts: Dict[str, int] = {}
        try:
            import docker

            self._client = docker.from_env() if not docker_host else docker.DockerClient(base_url=docker_host)
            # Quick connectivity check
            self._client.ping()
            logger.info("DockerOrchestrator connected to Docker daemon")
        except ImportError:
            logger.warning("Docker SDK not installed. Install with: pip install docker")
        except Exception as e:
            logger.warning(f"Could not connect to Docker daemon: {e}")
            self._client = None

    # ── Orchestrator properties ─────────────────────────────────────

    @property
    def available(self) -> bool:
        return self._client is not None

    @property
    def backend_name(self) -> str:
        return "docker"

    # ── Private helpers ─────────────────────────────────────────────

    def _build_container_env(
        self,
        *,
        run_id: str,
        store_url: str,
        env: Optional[Dict[str, str]],
        runner_config: Optional[Dict[str, Any]],
    ) -> Dict[str, str]:
        """Build the environment dict for an agent container."""
        import os as _os

        container_env: Dict[str, str] = {
            "MANTIS_RUN_ID": run_id,
            "MANTIS_STORE_URL": store_url,
        }

        # Forward Insight tracing config from worker env unconditionally
        for _key in (
            "INSIGHT_HOST",
            "INSIGHT_PUBLIC_KEY",
            "INSIGHT_SECRET_KEY",
        ):
            _val = _os.environ.get(_key)
            if _val:
                if _key == "INSIGHT_HOST" and ("localhost" in _val or "127.0.0.1" in _val):
                    import platform as _plat

                    _access = "host.docker.internal" if _plat.system() == "Darwin" else _val
                    _val = _val.replace("localhost", _access).replace("127.0.0.1", _access)
                container_env[_key] = _val

        # Fallback: forward LLM API keys from worker env ONLY if not already
        # provided via ProjectSecret (caller-provided env takes precedence).
        for _key in (
            "ANTHROPIC_API_KEY",
            "OPENAI_API_KEY",
            "AZURE_API_KEY",
            "GOOGLE_API_KEY",
            "COHERE_API_KEY",
        ):
            if _key not in container_env:
                _val = _os.environ.get(_key)
                if _val:
                    container_env[_key] = _val

        # Merge caller-provided env (takes precedence)
        if env:
            container_env.update(env)

        # Inject runner config as env vars (read by LitAgentRunner)
        if runner_config:
            if runner_config.get("max_rollouts") is not None:
                container_env["MANTIS_MAX_ROLLOUTS"] = str(runner_config["max_rollouts"])
            if runner_config.get("poll_interval") is not None:
                container_env["MANTIS_POLL_INTERVAL"] = str(runner_config["poll_interval"])
            if runner_config.get("n_runners") is not None:
                container_env["MANTIS_N_RUNNERS"] = str(runner_config["n_runners"])

        return container_env

    def _detect_network(self) -> str:
        """Auto-detect the Docker network to attach agent containers to."""
        if not self._client:
            return "bridge"
        try:
            networks = [n.name for n in self._client.networks.list()]
            for candidate in (
                "langfuse-network",
                "insight_default",
                "mantis_default",
            ):
                if candidate in networks:
                    logger.info(f"Auto-detected Docker network: {candidate}")
                    return candidate
        except Exception:
            pass
        return "bridge"

    def _build_docker_kwargs(
        self,
        *,
        container_env: Dict[str, str],
        container_labels: Dict[str, str],
        network: str,
        resource_limits: Optional[Dict[str, Any]],
        name: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Build the kwargs dict for ``docker.containers.run()``."""
        docker_kwargs: Dict[str, Any] = {
            "detach": True,
            "network": network,
            "environment": container_env,
            "labels": container_labels,
            "auto_remove": False,
        }

        if name:
            docker_kwargs["name"] = name

        if resource_limits:
            if resource_limits.get("memory"):
                docker_kwargs["mem_limit"] = resource_limits["memory"]
            if resource_limits.get("cpu"):
                docker_kwargs["nano_cpus"] = int(float(resource_limits["cpu"]) * 1e9)
            gpu = resource_limits.get("gpu")
            if gpu and gpu is not False:
                try:
                    import docker as _docker

                    gpu_count = -1 if gpu is True else int(gpu)
                    docker_kwargs["device_requests"] = [
                        _docker.types.DeviceRequest(count=gpu_count, capabilities=[["gpu"]])
                    ]
                except (ImportError, Exception) as e:
                    logger.warning(f"GPU device request failed: {e}")

        return docker_kwargs

    def _start_one_container(
        self,
        *,
        image: str,
        run_id: str,
        index: int,
        container_env: Dict[str, str],
        network: str,
        labels: Optional[Dict[str, str]],
        resource_limits: Optional[Dict[str, Any]],
    ) -> Optional[str]:
        """Start a single agent container. Returns container ID or None."""
        if not self._client:
            return None

        name = f"mantis-{run_id[:8]}-agent-{index}"
        container_labels = {
            "mantis.run_id": run_id,
            "mantis.component": "agent",
            "mantis.agent_index": str(index),
            **(labels or {}),
        }

        docker_kwargs = self._build_docker_kwargs(
            container_env=container_env,
            container_labels=container_labels,
            network=network,
            resource_limits=resource_limits,
            name=name,
        )

        try:
            container = self._client.containers.run(image, **docker_kwargs)
            logger.info(f"Started container {name} ({container.short_id}) " f"from {image} for run {run_id}")
            return container.id
        except Exception as e:
            logger.error(f"Failed to start container {name}: {e}")
            return None

    # ── Orchestrator interface ──────────────────────────────────────

    def start(
        self,
        *,
        run_id: str,
        image: str,
        store_url: str,
        count: int = 0,
        lifecycle: ContainerLifecycle = ContainerLifecycle.PERSISTENT,
        env: Optional[Dict[str, str]] = None,
        labels: Optional[Dict[str, str]] = None,
        runner_config: Optional[Dict[str, Any]] = None,
        resource_limits: Optional[Dict[str, Any]] = None,
        store: Optional[Any] = None,
        n_runners_per_unit: int = 1,
        max_units: int = 20,
        trainer_resources: Optional[Dict[str, Any]] = None,
        services: Optional[Dict[str, Any]] = None,
        algorithm_config: Optional[Dict[str, Any]] = None,
    ) -> ComputeHandle:
        if not self._client:
            raise RuntimeError("Docker not available, cannot start containers")

        container_env = self._build_container_env(
            run_id=run_id,
            store_url=store_url,
            env=env,
            runner_config=runner_config,
        )
        network = self._detect_network()

        # Store stop timeout for use in stop()
        stop_timeout = (resource_limits or {}).get("stop_timeout")
        if stop_timeout:
            self._stop_timeouts[run_id] = int(stop_timeout)

        identifiers: Dict[str, str] = {}
        for i in range(count):
            cid = self._start_one_container(
                image=image,
                run_id=run_id,
                index=i,
                container_env=container_env,
                network=network,
                labels=labels,
                resource_limits=resource_limits,
            )
            if cid is None:
                # Clean up any containers that were already started
                for key, started_id in identifiers.items():
                    self._stop_container(started_id)
                raise RuntimeError(f"Failed to start container agent-{i} " f"({i}/{count} started successfully)")
            identifiers[f"agent-{i}"] = cid

        if count > 0:
            logger.info(f"Started {count} agent container(s) for run {run_id} " f"(lifecycle={lifecycle.value})")

        handle = ComputeHandle(
            run_id=run_id,
            backend=self.backend_name,
            identifiers=identifiers,
            metadata={
                "lifecycle": lifecycle.value,
                "desired_count": count,
                "image": image,
                "store_url": store_url,
                "network": network,
                "env": env or {},
                "runner_config": runner_config or {},
                "resource_limits": resource_limits or {},
                "labels": labels or {},
            },
        )

        # If a store is provided, set up reactive scaling automatically.
        if store is not None:
            self._setup_reactive_listener(
                handle,
                store,
                n_runners_per_unit,
                max_units,
            )

        return handle

    def stop(self, handle: ComputeHandle, timeout: int = 30) -> None:
        if not self._client:
            return

        # Remove reactive listener first so no more events fire.
        self._teardown_reactive_listener(handle)

        stopped = 0
        for key, cid in list(handle.identifiers.items()):
            self._stop_container(cid)
            stopped += 1

        if stopped:
            logger.info(f"Killed and removed {stopped} container(s) for " f"run {handle.run_id}")

    def status(self, handle: ComputeHandle) -> ComputeStatus:
        if not self._client:
            return ComputeStatus.UNKNOWN

        statuses = []
        for key, cid in handle.identifiers.items():
            s = self._get_container_status(cid)
            statuses.append(s)

        if not statuses:
            return ComputeStatus.UNKNOWN
        if all(s == "running" for s in statuses):
            return ComputeStatus.RUNNING
        if any(s in ("exited", "dead") for s in statuses):
            return ComputeStatus.FAILED
        if all(s in ("created", "restarting") for s in statuses):
            return ComputeStatus.PENDING
        # Mixed state (some running, some starting) -- still consider running
        if any(s == "running" for s in statuses):
            return ComputeStatus.RUNNING
        return ComputeStatus.UNKNOWN

    def wait_ready(
        self,
        handle: ComputeHandle,
        timeout: float = 120.0,
    ) -> None:
        if not self._client:
            raise RuntimeError("Docker not available")

        deadline = _time.monotonic() + timeout
        poll_interval = 2.0
        target_count = len(handle.identifiers)

        while _time.monotonic() < deadline:
            ready = 0
            for key, cid in handle.identifiers.items():
                s = self._get_container_status(cid)
                if s == "running":
                    ready += 1
                elif s in ("exited", "dead"):
                    raise RuntimeError(f"Container {key} ({cid[:12]}) exited before " f"becoming ready")

            if ready >= target_count:
                logger.info(f"All {target_count} container(s) ready for " f"run {handle.run_id}")
                return

            logger.debug(
                f"Waiting for containers: {ready}/{target_count} ready "
                f"({timeout - (_time.monotonic() - (deadline - timeout)):.0f}s remaining)"
            )
            _time.sleep(poll_interval)

        raise RuntimeError(
            f"Timeout ({timeout}s) waiting for {target_count} container(s) " f"to become ready for run {handle.run_id}"
        )

    def scale(self, handle: ComputeHandle, count: int) -> None:
        if not self._client:
            raise RuntimeError("Docker not available")

        current = len(handle.identifiers)
        if count == current:
            return

        meta = handle.metadata

        if count > current:
            # Scale up: start additional containers
            container_env = self._build_container_env(
                run_id=handle.run_id,
                store_url=meta["store_url"],
                env=meta.get("env"),
                runner_config=meta.get("runner_config"),
            )
            network = meta.get("network", self._detect_network())

            for i in range(current, count):
                cid = self._start_one_container(
                    image=meta["image"],
                    run_id=handle.run_id,
                    index=i,
                    container_env=container_env,
                    network=network,
                    labels=meta.get("labels"),
                    resource_limits=meta.get("resource_limits"),
                )
                if cid is None:
                    logger.error(f"Failed to scale up: could not start agent-{i}")
                    break
                handle.identifiers[f"agent-{i}"] = cid

            logger.info(f"Scaled up from {current} to {len(handle.identifiers)} " f"for run {handle.run_id}")
        else:
            # Scale down: stop highest-indexed containers
            keys_sorted = sorted(
                handle.identifiers.keys(),
                key=lambda k: int(k.split("-")[-1]),
                reverse=True,
            )
            to_remove = keys_sorted[: current - count]
            for key in to_remove:
                cid = handle.identifiers.pop(key)
                self._stop_container(cid)

            logger.info(f"Scaled down from {current} to {len(handle.identifiers)} " f"for run {handle.run_id}")

        handle.metadata["desired_count"] = count

    def replace_member(
        self,
        handle: ComputeHandle,
        member_key: str,
    ) -> Optional[str]:
        if not self._client:
            return None

        old_cid = handle.identifiers.get(member_key)
        if old_cid:
            self._stop_container(old_cid)

        meta = handle.metadata
        index = int(member_key.split("-")[-1])

        container_env = self._build_container_env(
            run_id=handle.run_id,
            store_url=meta["store_url"],
            env=meta.get("env"),
            runner_config=meta.get("runner_config"),
        )
        network = meta.get("network", self._detect_network())

        new_cid = self._start_one_container(
            image=meta["image"],
            run_id=handle.run_id,
            index=index,
            container_env=container_env,
            network=network,
            labels=meta.get("labels"),
            resource_limits=meta.get("resource_limits"),
        )

        if new_cid:
            handle.identifiers[member_key] = new_cid
            logger.info(
                f"Replaced {member_key} ({old_cid[:12] if old_cid else '?'}) "
                f"with {new_cid[:12]} for run {handle.run_id}"
            )
        else:
            logger.error(f"Failed to replace {member_key} for run {handle.run_id}")

        return new_cid

    def members(self, handle: ComputeHandle) -> List[Dict[str, Any]]:
        if not self._client:
            return []

        result = []
        for key, cid in sorted(handle.identifiers.items()):
            index = int(key.split("-")[-1])
            status = self._get_container_status(cid) or "unknown"
            result.append(
                {
                    "key": key,
                    "id": cid,
                    "short_id": cid[:12],
                    "index": index,
                    "status": status,
                    "name": f"mantis-{handle.run_id[:8]}-agent-{index}",
                }
            )
        return result

    def get_logs(self, handle: ComputeHandle, tail: int = 100) -> str:
        if not self._client:
            return ""

        parts: List[str] = []
        for key, cid in sorted(handle.identifiers.items()):
            try:
                container = self._client.containers.get(cid)
                logs = container.logs(tail=tail, timestamps=False)
                if isinstance(logs, bytes):
                    logs = logs.decode("utf-8", errors="replace")
                parts.append(f"=== {key} ({cid[:12]}) ===\n{logs}")
            except Exception as e:
                parts.append(f"=== {key} ({cid[:12]}) === [error: {e}]")
        return "\n".join(parts)

    # ── Low-level container helpers ─────────────────────────────────

    def _stop_container(self, container_id: str) -> None:
        """Kill and remove a container immediately."""
        if not self._client:
            return
        try:
            container = self._client.containers.get(container_id)
            try:
                container.kill()
            except Exception:
                pass  # Already stopped / exited
            container.remove(force=True)
            logger.debug(f"Killed and removed container {container_id[:12]}")
        except Exception as e:
            logger.warning(f"Could not kill/remove container {container_id[:12]}: {e}")

    def _get_container_status(self, container_id: str) -> Optional[str]:
        """Get container status (running, exited, etc.)."""
        if not self._client:
            return None
        try:
            container = self._client.containers.get(container_id)
            return container.status
        except Exception as e:
            logger.debug(f"Could not get container status: {e}")
            return None

    # ── Legacy helpers (used by stop_all_for_run fallback) ──────────

    def list_run_containers(self, run_id: str) -> List[Dict[str, Any]]:
        """List all containers associated with a run (by label)."""
        if not self._client:
            return []
        try:
            containers = self._client.containers.list(
                all=True,
                filters={"label": f"mantis.run_id={run_id}"},
            )
            return [
                {
                    "id": c.id,
                    "short_id": c.short_id,
                    "name": c.name,
                    "status": c.status,
                    "image": str(c.image),
                }
                for c in containers
            ]
        except Exception as e:
            logger.debug(f"Could not list containers for run {run_id}: {e}")
            return []

    def stop_all_for_run(self, run_id: str) -> int:
        """Kill all containers associated with a run (label-based fallback)."""
        containers = self.list_run_containers(run_id)
        stopped = 0
        for c in containers:
            self._stop_container(c["id"])
            stopped += 1
        return stopped

    # ── Reactive scaling / ephemeral recycling ───────────────────

    def _setup_reactive_listener(
        self,
        handle: ComputeHandle,
        store: Any,
        n_runners_per_unit: int,
        max_units: int,
    ) -> None:
        """Register a store listener for on-demand scaling and ephemeral replacement.

        Called automatically from ``start()`` when a ``store`` is provided.
        """
        import math

        lifecycle_str = handle.metadata.get("lifecycle", "persistent")
        is_ephemeral = lifecycle_str == ContainerLifecycle.EPHEMERAL.value
        orch = self

        class _ReactiveListener:
            """Store listener that starts/replaces containers on demand."""

            @property
            def capabilities(self) -> Dict[str, bool]:
                return {}

            def otlp_traces_endpoint(self) -> Optional[str]:
                return None

            def get_otlp_headers(self) -> Dict[str, str]:
                return {}

            async def on_rollout_created(self, rollout: Any) -> None:
                current_units = len(handle.identifiers)
                # Scale to one more container than we currently have,
                # but only if all existing slots are occupied.
                # Don't call wait_ready() here -- it blocks the event
                # loop and causes cascading scale-ups.  The container
                # will become ready asynchronously; the runner inside
                # polls the store and picks up work once healthy.
                if current_units < max_units:
                    needed = current_units + 1
                    try:
                        orch.scale(handle, needed)
                        logger.info(
                            f"Reactive: started container agent-{current_units} "
                            f"for run {handle.run_id} "
                            f"({needed} total)"
                        )
                    except Exception as e:
                        logger.warning(f"Reactive scale-up failed for " f"run {handle.run_id}: {e}")

            async def on_rollout_updated(self, rollout: Any) -> None:
                if not is_ephemeral:
                    return
                status = getattr(rollout, "status", None)
                if status not in ("succeeded", "failed"):
                    return

                # Replace the first container that is no longer running
                # (exited after processing its rollout).
                for key, cid in list(handle.identifiers.items()):
                    cstatus = orch._get_container_status(cid)
                    if cstatus in ("exited", "dead", None):
                        try:
                            new_cid = orch.replace_member(handle, key)
                            if new_cid:
                                logger.info(
                                    f"Ephemeral recycle: replaced {key} "
                                    f"({cid[:12]}) with {new_cid[:12]} "
                                    f"for run {handle.run_id}"
                                )
                        except Exception as e:
                            logger.warning(f"Ephemeral recycle failed for {key}: {e}")
                        break

        listener = _ReactiveListener()
        if not hasattr(self, "_reactive_listeners"):
            self._reactive_listeners: Dict[str, Any] = {}
        self._reactive_listeners[handle.run_id] = (listener, store)

        store.add_listener(listener)
        logger.info(
            f"Reactive listener registered for run {handle.run_id} "
            f"(ephemeral={is_ephemeral}, n_runners={n_runners_per_unit}, "
            f"max_units={max_units})"
        )

    def _teardown_reactive_listener(self, handle: ComputeHandle) -> None:
        """Remove the reactive listener from the store."""
        listeners = getattr(self, "_reactive_listeners", {})
        entry = listeners.pop(handle.run_id, None)
        if entry:
            listener, store = entry
            store.remove_listener(listener)
            logger.debug(f"Reactive listener removed for run {handle.run_id}")
