"""Eject policies — early termination conditions for agent runs."""

from __future__ import annotations

import time
from typing import Any, Callable

from .models import EjectDecision


class EjectPolicy:
    """Base class for eject policies. Subclass and override :meth:`check`."""

    def check(
        self,
        *,
        elapsed_s: float,
        budget_status: Any,  # BudgetStatus
        events: list[dict],
        context: dict,
    ) -> EjectDecision:
        """Evaluate whether the agent should be terminated.

        Args:
            elapsed_s: Seconds since the run started.
            budget_status: Current :class:`~agent_preflight.models.BudgetStatus`.
            events: All events captured so far.
            context: Arbitrary run-specific state (e.g. tool call count, progress flags).

        Returns:
            :class:`EjectDecision`.
        """
        return EjectDecision(should_eject=False)


class BudgetExceededEject(EjectPolicy):
    """Eject when measured spend reaches the budget limit.

    This is the failsafe: OpenRouter sub-keys have a soft spend cap, but the
    harness enforces a hard kill via this policy.
    """

    def check(self, *, budget_status: Any, **_kwargs: Any) -> EjectDecision:
        if budget_status.exceeded:
            return EjectDecision(
                should_eject=True,
                reason=f"budget exhausted (${budget_status.used:.3f} >= ${budget_status.limit:.2f})",
            )
        return EjectDecision(should_eject=False)


class IdleTimeoutEject(EjectPolicy):
    """Eject if no new events have arrived for ``timeout_s`` seconds.

    Use this to kill containers that are stuck (e.g. waiting on a hung tool call).
    """

    def __init__(self, timeout_s: float = 300.0) -> None:
        self.timeout_s = timeout_s
        self._last_event_count: int = 0
        self._last_event_ts: float = time.time()

    def check(self, *, events: list[dict], **_kwargs: Any) -> EjectDecision:
        if len(events) > self._last_event_count:
            self._last_event_count = len(events)
            self._last_event_ts = time.time()

        idle = time.time() - self._last_event_ts
        if idle >= self.timeout_s:
            return EjectDecision(
                should_eject=True,
                reason=f"no events for {idle:.0f}s (timeout={self.timeout_s:.0f}s)",
            )
        return EjectDecision(should_eject=False)


class BudgetFractionEject(EjectPolicy):
    """Eject once the agent has spent a given fraction of its budget with no progress.

    The ``progress_fn`` receives the current context dict and should return True
    if the agent has made meaningful progress. If it returns False at
    ``check_after_pct`` budget spent, the agent is ejected.

    Example — eject after 20% budget spent if no tool calls recorded::

        BudgetFractionEject(
            check_after_pct=0.20,
            progress_fn=lambda ctx: ctx.get("tool_calls", 0) > 0,
        )
    """

    def __init__(
        self,
        check_after_pct: float = 0.20,
        progress_fn: Callable[[dict], bool] | None = None,
    ) -> None:
        self.check_after_pct = check_after_pct
        self.progress_fn = progress_fn or (lambda ctx: True)
        self._fired: bool = False

    def check(self, *, budget_status: Any, context: dict, **_kwargs: Any) -> EjectDecision:
        if self._fired:
            return EjectDecision(should_eject=False)
        if budget_status.pct_spent >= self.check_after_pct:
            self._fired = True
            if not self.progress_fn(context):
                return EjectDecision(
                    should_eject=True,
                    reason=(
                        f"no progress at {budget_status.pct_spent*100:.0f}% budget spent "
                        f"(${budget_status.used:.3f} / ${budget_status.limit:.2f})"
                    ),
                )
        return EjectDecision(should_eject=False)


class CompositeEject(EjectPolicy):
    """Combine multiple eject policies with OR logic (first to trigger wins)."""

    def __init__(self, policies: list[EjectPolicy]) -> None:
        self.policies = policies

    def check(self, **kwargs: Any) -> EjectDecision:
        for policy in self.policies:
            decision = policy.check(**kwargs)
            if decision.should_eject:
                return decision
        return EjectDecision(should_eject=False)


def default_eject_policies(idle_timeout_s: float = 300.0) -> CompositeEject:
    """Return a sensible default set: budget exhausted + idle timeout.

    Args:
        idle_timeout_s: Seconds of silence before killing. Default 5 minutes.
    """
    return CompositeEject([
        BudgetExceededEject(),
        IdleTimeoutEject(timeout_s=idle_timeout_s),
    ])
