import logging
from typing import Any, Iterable, Optional, Tuple, Dict
import json
import time
import contextvars

from litellm.exceptions import APIError, BadRequestError, NotFoundError, OpenAIError, ServiceUnavailableError

from api.models import (
    BrowserUseAgentTask,
    PersistentAgent,
    PersistentAgentMessage,
    PersistentAgentStep,
    EvalRunTask,
    EvalRun,
    CommsAllowlistEntry
)
from api.agent.comms.message_service import inject_internal_web_message
from api.agent.core.llm_utils import LiteLLMResponseError, run_completion
from api.agent.events import AgentEventType, get_agent_event_stream
from api.evals.realtime import broadcast_task_update, broadcast_run_update
from api.evals.metrics import aggregate_run_metrics
from config.redis_client import get_redis_client
from django.utils import timezone

logger = logging.getLogger(__name__)

_JUDGE_RETRY_DELAYS_SECONDS = (1.0, 3.0, 6.0, 10.0)
_JUDGE_RETRYABLE_ERRORS = (
    AttributeError,
    APIError,
    BadRequestError,
    IndexError,
    json.JSONDecodeError,
    KeyError,
    LiteLLMResponseError,
    NotFoundError,
    OpenAIError,
    ServiceUnavailableError,
    TypeError,
    ValueError,
)


def _is_structured_judge_grammar_error(exc: Exception) -> bool:
    message = str(exc).lower()
    return (
        "structural_tag grammar" in message
        or "failed to compile structural" in message
        or "tool_choice" in message
        or "submit_judgment" in message
    )


def _preview_text(value: Any, *, limit: int = 1200) -> str:
    text = str(value)
    if len(text) <= limit:
        return text
    return text[: limit - 3] + "..."


def _serialize_debug_artifact(value: Any, *, depth: int = 0) -> Any:
    if depth > 3:
        return _preview_text(value, limit=300)
    if value is None or isinstance(value, (bool, int, float)):
        return value
    if isinstance(value, str):
        return _preview_text(value)
    if isinstance(value, PersistentAgentMessage):
        return {
            "type": "message",
            "id": str(value.id),
            "is_outbound": value.is_outbound,
            "timestamp": value.timestamp.isoformat() if value.timestamp else None,
            "body_preview": _preview_text(value.body or "", limit=1200),
        }
    if isinstance(value, PersistentAgentStep):
        return {
            "type": "step",
            "id": str(value.id),
            "created_at": value.created_at.isoformat() if value.created_at else None,
            "description": _preview_text(value.description or "", limit=600),
        }
    if isinstance(value, BrowserUseAgentTask):
        return {
            "type": "browser_task",
            "id": str(value.id),
            "status": value.status,
            "created_at": value.created_at.isoformat() if value.created_at else None,
        }
    if isinstance(value, dict):
        return {
            str(key): _serialize_debug_artifact(item, depth=depth + 1)
            for key, item in value.items()
            if not str(key).lower().endswith(("key", "token", "secret", "password"))
        }
    if isinstance(value, (list, tuple, set)):
        return [
            _serialize_debug_artifact(item, depth=depth + 1)
            for item in list(value)[:25]
        ]
    if hasattr(value, "pk") and hasattr(value, "_meta"):
        return {
            "type": value._meta.label_lower,
            "id": str(value.pk),
        }
    return _preview_text(value)


def _sanitize_debug_artifacts(artifacts: Dict[str, Any]) -> dict[str, Any]:
    sanitized = {}
    for key, value in (artifacts or {}).items():
        sanitized[str(key)] = _serialize_debug_artifact(value)
    return sanitized


def _response_message_content(response: Any) -> str:
    choices = response.get("choices") if isinstance(response, dict) else getattr(response, "choices", None)
    if not choices:
        raise ValueError("LLM judge JSON fallback returned no choices.")
    first_choice = choices[0]
    message = first_choice.get("message") if isinstance(first_choice, dict) else getattr(first_choice, "message", None)
    if message is None:
        raise ValueError("LLM judge JSON fallback returned no message.")
    content = message.get("content") if isinstance(message, dict) else getattr(message, "content", None)
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for part in content:
            if isinstance(part, str):
                parts.append(part)
                continue
            if isinstance(part, dict) and isinstance(part.get("text"), str):
                parts.append(part["text"])
        return "\n".join(parts)
    raise ValueError("LLM judge JSON fallback returned no text content.")


def _strip_json_code_fence(text: str) -> str:
    stripped = text.strip()
    if not stripped.startswith("```"):
        return stripped
    lines = stripped.splitlines()
    if len(lines) >= 3 and lines[-1].strip() == "```":
        return "\n".join(lines[1:-1]).strip()
    return stripped


def _extract_json_object_from_text(text: str) -> dict[str, Any]:
    stripped = _strip_json_code_fence(text)
    if not stripped:
        raise ValueError("LLM judge JSON fallback returned empty content.")

    decoder = json.JSONDecoder()
    starts = [0] + [idx for idx, char in enumerate(stripped) if char == "{" and idx != 0]
    for start in starts:
        try:
            payload, _end = decoder.raw_decode(stripped[start:])
        except json.JSONDecodeError:
            continue
        if isinstance(payload, dict):
            return payload
    raise ValueError("LLM judge JSON fallback did not return a JSON object.")


_current_eval_run_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
    "eval_run_id", default=None
)

_current_eval_routing_profile: contextvars.ContextVar[Any] = contextvars.ContextVar(
    "eval_routing_profile", default=None
)


def set_current_eval_run_id(run_id: str | None) -> None:
    _current_eval_run_id.set(run_id)


def get_current_eval_run_id() -> str | None:
    return _current_eval_run_id.get()


def set_current_eval_routing_profile(profile: Any) -> None:
    """Set the routing profile for the current eval context."""
    _current_eval_routing_profile.set(profile)


def get_current_eval_routing_profile() -> Any:
    """Get the routing profile for the current eval context, or None."""
    return _current_eval_routing_profile.get()


def get_eval_routing_profile_for_current_run() -> Any:
    """Resolve the eval routing profile from context, then persisted run state."""
    routing_profile = get_current_eval_routing_profile()
    if routing_profile is not None:
        return routing_profile

    run_id = get_current_eval_run_id()
    if not run_id:
        return None

    try:
        run = (
            EvalRun.objects
            .select_related("llm_routing_profile", "suite_run__llm_routing_profile")
            .get(id=run_id)
        )
    except EvalRun.DoesNotExist:
        return None

    if run.llm_routing_profile_id:
        return run.llm_routing_profile
    if run.suite_run_id and run.suite_run.llm_routing_profile_id:
        return run.suite_run.llm_routing_profile
    return None

class AgentEventListener:
    """
    Lightweight, event-driven listener that reads the agent event stream.
    Designed to avoid polling and to tolerate events that were emitted just
    before the listener started (by filtering on start_time).
    """
    def __init__(self, agent_id: str, *, start_time: Optional[float] = None):
        self.agent_id = str(agent_id)
        self.stream_key = get_agent_event_stream(agent_id)
        self.start_time = start_time or time.time()
        self.redis = get_redis_client()
        self.last_id = "0-0"

    def __enter__(self) -> "AgentEventListener":
        return self

    def __exit__(self, exc_type, exc_value, traceback) -> bool:
        # No resources to release; propagate exceptions if any
        return False

    def wait_for(
        self,
        event_type: AgentEventType | str,
        timeout: int = 30,
    ) -> Optional[Dict[str, Any]]:
        """
        Block on the Redis stream for a matching event or until timeout.
        Returns the decoded event dict or None on timeout/error.
        """
        target_type = event_type.value if isinstance(event_type, AgentEventType) else str(event_type)
        deadline = time.time() + timeout

        while time.time() < deadline:
            remaining_ms = max(1, int((deadline - time.time()) * 1000))
            try:
                messages = self.redis.xread(
                    {self.stream_key: self.last_id},
                    count=50,
                    block=remaining_ms,
                )
            except Exception:
                logger.warning("Failed to read agent event stream for %s", self.agent_id, exc_info=True)
                return None

            if not messages:
                continue

            for _stream, entries in messages:
                for entry_id, raw_fields in entries:
                    self.last_id = entry_id
                    raw = raw_fields.get("data") or raw_fields.get(b"data")
                    if raw is None:
                        continue
                    try:
                        if isinstance(raw, bytes):
                            event = json.loads(raw.decode("utf-8"))
                        elif isinstance(raw, str):
                            event = json.loads(raw)
                        elif isinstance(raw, dict):
                            event = raw
                        else:
                            continue
                    except Exception:
                        continue

                    try:
                        ts_val = float(event.get("timestamp", 0))
                    except Exception:
                        ts_val = 0.0

                    # Ignore stale events that occurred before we started listening
                    if ts_val < self.start_time:
                        continue

                    if event.get("type") == target_type:
                        return event

        return None

class ScenarioExecutionTools:
    """
    Tools for scenarios to interact with the agent and record results.
    Intended to be used as a base class or mixin for EvalScenario.
    """

    def get_agent(self, agent_id: str) -> PersistentAgent:
        return PersistentAgent.objects.get(id=agent_id)
    
    def get_run(self, run_id: str) -> EvalRun:
        return EvalRun.objects.get(id=run_id)

    def inject_message(
        self,
        agent_id: str,
        body: str,
        sender_user_id: int = -999,
        attachments: Iterable[Any] = (),
        trigger_processing: bool = True,
        eval_run_id: str | None = None,
        mock_config: dict | None = None,
        eval_stop_policy: dict | None = None,
    ) -> PersistentAgentMessage:
        """
        Send a message to the agent as a web user.
        Automatically whitelists the sender to ensure the agent can reply.

        Args:
            mock_config: Optional dict mapping tool_name -> mock response.
                         Passed to Celery worker for eval mocking.
        """
        current_run_id = eval_run_id or get_current_eval_run_id()
        agent = PersistentAgent.objects.select_related("user").get(id=agent_id)
        resolved_sender_user_id = sender_user_id
        if current_run_id and sender_user_id == -999 and agent.user_id:
            resolved_sender_user_id = agent.user_id

        msg, _ = inject_internal_web_message(
            agent_id=agent_id,
            body=body,
            sender_user_id=resolved_sender_user_id,
            attachments=attachments,
            trigger_processing=False,  # handle processing explicitly below
            eval_run_id=current_run_id,
        )

        # Auto-whitelist the sender so the agent trusts this contact
        CommsAllowlistEntry.objects.get_or_create(
            agent_id=agent_id,
            channel=msg.from_endpoint.channel,
            address=msg.from_endpoint.address,
            defaults={
                "is_active": True,
            }
        )

        # Update agent's preferred contact to this new user so "Welcome" prompts target them
        agent.preferred_contact_endpoint = msg.from_endpoint
        agent.save(update_fields=["preferred_contact_endpoint"])

        if current_run_id and resolved_sender_user_id == agent.user_id:
            from api.services.web_sessions import start_web_session

            start_web_session(agent, agent.user, source="eval")

        if trigger_processing:
            self._dispatch_agent_processing(
                agent_id,
                eval_run_id=current_run_id,
                mock_config=mock_config,
                eval_stop_policy=eval_stop_policy,
            )

        return msg

    def trigger_processing(
        self,
        agent_id: str,
        *,
        eval_run_id: str | None = None,
        mock_config: dict | None = None,
        eval_stop_policy: dict | None = None,
    ) -> None:
        """
        Manually trigger the agent's event processing loop.
        """
        current_run_id = eval_run_id or get_current_eval_run_id()
        self._dispatch_agent_processing(
            agent_id,
            eval_run_id=current_run_id,
            mock_config=mock_config,
            eval_stop_policy=eval_stop_policy,
        )

    def _dispatch_agent_processing(
        self,
        agent_id: str,
        *,
        eval_run_id: str | None = None,
        mock_config: dict | None = None,
        eval_stop_policy: dict | None = None,
    ) -> None:
        """
        Start agent processing for scenario execution.

        Eval runs execute the agent inline so a single Celery worker does not deadlock by
        waiting on a process_agent_events task queued behind the active run_eval_task.
        """
        try:
            from api.agent.tasks import process_agent_events_task

            kwargs = {
                "eval_run_id": eval_run_id,
                "mock_config": mock_config,
            }
            if eval_stop_policy is not None:
                kwargs["eval_stop_policy"] = eval_stop_policy

            if eval_run_id:
                process_agent_events_task.apply(
                    args=(str(agent_id),),
                    kwargs=kwargs,
                    throw=True,
                )
            else:
                process_agent_events_task.delay(
                    str(agent_id),
                    **kwargs,
                )
        except Exception:
            logger.exception("Failed to trigger processing for agent %s", agent_id)
            if eval_run_id:
                raise

    def agent_event_listener(self, agent_id: str, *, start_time: Optional[float] = None) -> AgentEventListener:
        """
        Convenience helper to create an AgentEventListener with a start timestamp.
        """
        return AgentEventListener(agent_id, start_time=start_time)

    def record_task_result(
        self,
        run_id: str,
        task_sequence: Optional[int],
        status: str,
        observed_summary: str = "",
        expected_summary: str = "",
        artifacts: Dict[str, Any] = None,
        task_name: Optional[str] = None
    ) -> EvalRunTask:
        """
        Update or create a task result record.
        """
        artifacts = artifacts or {}
        
        if task_name:
            task_obj = EvalRunTask.objects.get(run_id=run_id, name=task_name)
        elif task_sequence is not None:
            task_obj, created = EvalRunTask.objects.get_or_create(
                run_id=run_id,
                sequence=task_sequence,
                defaults={
                    "name": f"Task {task_sequence}",
                    "assertion_type": "manual"
                }
            )
        else:
            raise ValueError("Must provide either task_sequence or task_name")
        
        task_obj.status = status
        if observed_summary:
            task_obj.observed_summary = observed_summary
        if expected_summary:
            task_obj.expected_summary = expected_summary

        now = timezone.now()
        if task_obj.started_at is None and status in (
            EvalRunTask.Status.RUNNING,
            EvalRunTask.Status.PASSED,
            EvalRunTask.Status.FAILED,
            EvalRunTask.Status.ERRORED,
            EvalRunTask.Status.SKIPPED,
        ):
            task_obj.started_at = now

        if status in (
            EvalRunTask.Status.PASSED,
            EvalRunTask.Status.FAILED,
            EvalRunTask.Status.ERRORED,
            EvalRunTask.Status.SKIPPED,
        ):
            task_obj.finished_at = now
            
        # Link artifacts if provided
        if "message" in artifacts:
            task_obj.first_message = artifacts["message"]
        if "step" in artifacts:
            task_obj.first_step = artifacts["step"]
        if "browser_task" in artifacts:
            task_obj.first_browser_task = artifacts["browser_task"]
        debug_artifacts = _sanitize_debug_artifacts(artifacts)
        if debug_artifacts:
            task_obj.debug_artifacts = {
                **(task_obj.debug_artifacts or {}),
                **debug_artifacts,
            }
            
        task_obj.save()

        # Attempt to aggregate cost/usage metrics for this task and its parent run.
        # We call aggregate_run_metrics, which will:
        # 1. Sum total costs for the run from all AgentCompletions/Steps.
        # 2. Re-distribute those costs to tasks based on time windows.
        try:
            aggregate_run_metrics(task_obj.run)
            broadcast_run_update(task_obj.run)
        except Exception:
            logger.error("Failed to aggregate metrics during record_task_result", exc_info=True)

        try:
            broadcast_task_update(task_obj)
        except Exception:
            logger.debug("Broadcast task update failed", exc_info=True)

        return task_obj

    def llm_judge(
        self,
        question: str,
        context: str,
        options: Iterable[str] = ("Yes", "No"),
        model: Optional[str] = None,
        params: Optional[Dict[str, Any]] = None,
    ) -> Tuple[str, str]:
        """
        Ask an LLM to judge a context based on a question and a set of options.
        Uses tool calling to ensure structured output. Uses the eval routing profile
        if one is set, otherwise falls back to the active profile or legacy config.

        Args:
            question: The specific question to answer.
            context: The context text to evaluate.
            options: A list of valid answer options (default: ["Yes", "No"]).
            model: Optional LLM model to use. If omitted, the configured default is used.
            params: Optional LLM parameters. If omitted, the configured default params are used.

        Returns:
            A tuple of (choice, reasoning). choice will be one of the strings in `options`.
        """
        from api.agent.core.llm_config import get_llm_config_with_failover, LLMNotConfiguredError

        options_list = list(options)

        tool_definition = {
            "type": "function",
            "function": {
                "name": "submit_judgment",
                "description": "Submit the final judgment for the evaluation.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "choice": {
                            "type": "string",
                            "enum": options_list,
                            "description": "The selected answer from the valid options."
                        },
                        "reasoning": {
                            "type": "string",
                            "description": "Concise justification for the choice."
                        }
                    },
                    "required": ["choice", "reasoning"]
                }
            }
        }

        prompt = [
            {"role": "system", "content": "You are an impartial judge. Evaluate the context and answer the question by calling the `submit_judgment` tool."},
            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}\n\nValid Options: {', '.join(options_list)}"}
        ]
        judge_tool_choice = {"type": "function", "function": {"name": "submit_judgment"}}

        # If caller provided both model and params, use them directly
        if model is not None and params is not None:
            safe_params = dict(params)
            if safe_params.get("temperature") is None:
                safe_params["temperature"] = 0.0
            try:
                return self._run_judge_completion(
                    model=model,
                    prompt=prompt,
                    tool_definition=tool_definition,
                    tool_choice=judge_tool_choice,
                    params=safe_params,
                    options=options_list,
                )
            except _JUDGE_RETRYABLE_ERRORS as e:
                logger.error("LLM judge failed with explicit model %s: %s", model, e)
                return "Error", f"Exception during judgment: {str(e)}"

        # Use failover configs with routing profile support
        try:
            routing_profile = get_eval_routing_profile_for_current_run()
            failover_configs = get_llm_config_with_failover(routing_profile=routing_profile)
        except LLMNotConfiguredError as exc:
            logger.error("LLM judge missing configuration: %s", exc)
            return "Error", "No LLM configuration available for judgment."

        if not failover_configs:
            return "Error", "No LLM model available for judgment."

        last_error: Optional[Exception] = None
        for _provider, cfg_model, cfg_params in failover_configs:
            effective_model = model or cfg_model
            safe_params = dict(params or cfg_params or {})
            if safe_params.get("temperature") is None:
                safe_params["temperature"] = 0.0

            try:
                return self._run_judge_completion(
                    model=effective_model,
                    prompt=prompt,
                    tool_definition=tool_definition,
                    tool_choice=judge_tool_choice,
                    params=safe_params,
                    options=options_list,
                )
            except _JUDGE_RETRYABLE_ERRORS as e:
                last_error = e
                logger.warning("LLM judge failed with model %s: %s, trying next", effective_model, e)
                continue

        logger.error("LLM judge failed with all providers: %s", last_error)
        return "Error", f"Exception during judgment: {str(last_error)}"

    def _run_judge_completion(
        self,
        *,
        model: str,
        prompt: list[dict[str, str]],
        tool_definition: dict[str, Any],
        tool_choice: dict[str, Any],
        params: dict[str, Any],
        options: list[str],
    ) -> Tuple[str, str]:
        last_error: Exception | None = None
        for attempt, delay_seconds in enumerate((*_JUDGE_RETRY_DELAYS_SECONDS, None), start=1):
            try:
                response = run_completion(
                    model=model,
                    messages=prompt,
                    tools=[tool_definition],
                    tool_choice=tool_choice,
                    params=params,
                    drop_params=True,
                )
                choice, reasoning = self._extract_judgment(response)
                if choice not in options:
                    raise ValueError(f"LLM judge returned invalid choice {choice!r}: {reasoning}")
                return choice, reasoning
            except _JUDGE_RETRYABLE_ERRORS as exc:
                if _is_structured_judge_grammar_error(exc):
                    logger.warning(
                        "LLM judge structured tool call failed with provider grammar error for model %s; "
                        "retrying as strict JSON text.",
                        model,
                    )
                    return self._run_unstructured_judge_completion(
                        model=model,
                        prompt=prompt,
                        params=params,
                        options=options,
                    )
                last_error = exc
                if delay_seconds is None:
                    raise
                logger.warning(
                    "LLM judge returned retryable response error with model %s; retrying attempt %s: %s",
                    model,
                    attempt + 1,
                    exc,
                )
                time.sleep(delay_seconds)

        if last_error is not None:
            raise last_error
        raise ValueError("LLM judge failed without a captured error.")

    def _run_unstructured_judge_completion(
        self,
        *,
        model: str,
        prompt: list[dict[str, str]],
        params: dict[str, Any],
        options: list[str],
    ) -> Tuple[str, str]:
        fallback_prompt = [
            {
                "role": "system",
                "content": (
                    "You are an impartial judge. Evaluate the context and answer the question. "
                    "Return exactly one JSON object with keys choice and reasoning. Do not use markdown."
                ),
            },
            {
                "role": "user",
                "content": (
                    f"{prompt[1]['content']}\n\n"
                    "Return JSON only in this shape: "
                    '{"choice":"<one valid option>","reasoning":"<concise justification>"}'
                ),
            },
        ]
        last_error: Exception | None = None
        for attempt, delay_seconds in enumerate((*_JUDGE_RETRY_DELAYS_SECONDS, None), start=1):
            try:
                response = run_completion(
                    model=model,
                    messages=fallback_prompt,
                    params=params,
                    drop_params=True,
                )
                choice, reasoning = self._extract_unstructured_judgment(response)
                if choice not in options:
                    raise ValueError(f"LLM judge returned invalid choice {choice!r}: {reasoning}")
                return choice, f"Structured judge fallback used after provider grammar error. {reasoning}"
            except _JUDGE_RETRYABLE_ERRORS as exc:
                last_error = exc
                if delay_seconds is None:
                    raise
                logger.warning(
                    "LLM judge JSON fallback returned retryable response error with model %s; "
                    "retrying attempt %s: %s",
                    model,
                    attempt + 1,
                    exc,
                )
                time.sleep(delay_seconds)

        if last_error is not None:
            raise last_error
        raise ValueError("LLM judge JSON fallback failed without a captured error.")

    def _extract_judgment(self, response) -> Tuple[str, str]:
        """Extract judgment from LLM response tool calls."""
        tool_calls = response.choices[0].message.tool_calls
        if not tool_calls:
            return "Error", "LLM did not call the judgment tool."

        for tc in tool_calls:
            if tc.function.name == "submit_judgment":
                args = json.loads(tc.function.arguments)
                return args.get("choice"), args.get("reasoning")

        return "Error", "LLM did not call submit_judgment tool."

    def _extract_unstructured_judgment(self, response) -> Tuple[str, str]:
        payload = _extract_json_object_from_text(_response_message_content(response))
        choice = payload.get("choice")
        reasoning = payload.get("reasoning")
        if not isinstance(choice, str):
            raise ValueError("LLM judge JSON fallback did not return a string choice.")
        if not isinstance(reasoning, str):
            raise ValueError("LLM judge JSON fallback did not return string reasoning.")
        return choice, reasoning

    def wait_for_event(
        self,
        agent_id: str,
        event_type: str,
        timeout: int = 30,
        start_time: Optional[float] = None,
        listener: Optional[AgentEventListener] = None,
    ) -> Optional[Dict[str, Any]]:
        """
        Block until a specific event type is received for the agent using the event stream.
        Returns the full event payload if received, None if timeout.
        If a listener is provided, it will be used (and advanced) to avoid duplicate reads.
        """
        effective_listener = listener or AgentEventListener(agent_id, start_time=start_time)
        return effective_listener.wait_for(event_type, timeout=timeout)

    def wait_for_idle(self, agent_id: str, timeout: int = 60) -> bool:
        """
        Wait until the agent emits PROCESSING_COMPLETE with 0 outstanding tasks.
        Returns True if idle state reached, False if timeout.
        """
        listener = AgentEventListener(agent_id, start_time=time.time())
        deadline = time.time() + timeout
        remaining = timeout

        while remaining > 0:
            event = listener.wait_for(AgentEventType.PROCESSING_COMPLETE, timeout=int(remaining))
            if not event:
                return False
            outstanding = int((event.get("payload") or {}).get("outstanding_tasks", 0) or 0)
            if outstanding == 0:
                return True
            remaining = max(0, deadline - time.time())
        return False

    def wait_for_agent_idle(self, agent_id: str, timeout: int = 60):
        """
        Return a context manager that waits for the agent to become idle after the block executes.
        Usage:
            with self.wait_for_agent_idle(agent_id):
                self.inject_message(...)
        """
        return WaitForIdleContext(agent_id, timeout)

class WaitForIdleContext:
    """
    Context manager that subscribes to agent events BEFORE the action,
    then waits for the agent to go idle AFTER the action.
    Eliminates race conditions in eager/fast execution environments.
    """
    def __init__(self, agent_id: str, timeout: int = 60):
        self.agent_id = agent_id
        self.timeout = timeout
        self.listener: Optional[AgentEventListener] = None

    def __enter__(self):
        self.listener = AgentEventListener(self.agent_id, start_time=time.time())
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type:
            return False  # Propagate exception

        if not self.listener:
            return False

        deadline = time.time() + self.timeout
        remaining = self.timeout
        while remaining > 0:
            event = self.listener.wait_for(AgentEventType.PROCESSING_COMPLETE, timeout=int(remaining))
            if not event:
                break
            outstanding = int((event.get("payload") or {}).get("outstanding_tasks", 0) or 0)
            if outstanding == 0:
                return True  # Success
            remaining = max(0, deadline - time.time())

        logger.warning(f"Timeout waiting for agent {self.agent_id} to go idle.")
        return False # Do not suppress exceptions, but flow continues if no exception
