"""Tests for eject policies."""

import time

from agent_preflight.eject import (
    BudgetExceededEject,
    BudgetFractionEject,
    CompositeEject,
    IdleTimeoutEject,
    default_eject_policies,
)
from agent_preflight.models import BudgetStatus


def _budget(used: float, limit: float = 0.25) -> BudgetStatus:
    return BudgetStatus(used=used, limit=limit)


def _kwargs(budget_status, events=None, context=None, elapsed_s=60.0):
    return dict(
        budget_status=budget_status,
        events=events or [],
        context=context or {},
        elapsed_s=elapsed_s,
    )


# ---------------------------------------------------------------------------
# BudgetExceededEject
# ---------------------------------------------------------------------------


class TestBudgetExceededEject:
    def test_no_eject_under_budget(self):
        d = BudgetExceededEject().check(**_kwargs(_budget(0.10)))
        assert not d.should_eject

    def test_eject_at_limit(self):
        d = BudgetExceededEject().check(**_kwargs(_budget(0.25)))
        assert d.should_eject
        assert "exhausted" in d.reason

    def test_eject_over_limit(self):
        d = BudgetExceededEject().check(**_kwargs(_budget(0.30)))
        assert d.should_eject


# ---------------------------------------------------------------------------
# IdleTimeoutEject
# ---------------------------------------------------------------------------


class TestIdleTimeoutEject:
    def test_no_eject_with_events(self):
        policy = IdleTimeoutEject(timeout_s=1.0)
        d = policy.check(**_kwargs(_budget(0.10), events=[{"type": "x"}]))
        assert not d.should_eject

    def test_eject_after_timeout(self):
        policy = IdleTimeoutEject(timeout_s=0.01)
        # First check sets baseline
        policy.check(**_kwargs(_budget(0.10), events=[]))
        time.sleep(0.02)
        d = policy.check(**_kwargs(_budget(0.10), events=[]))
        assert d.should_eject
        assert "no events" in d.reason

    def test_reset_on_new_events(self):
        policy = IdleTimeoutEject(timeout_s=0.01)
        policy.check(**_kwargs(_budget(0.10), events=[]))
        time.sleep(0.02)
        # New event resets the timer
        d = policy.check(**_kwargs(_budget(0.10), events=[{"type": "x"}]))
        assert not d.should_eject


# ---------------------------------------------------------------------------
# BudgetFractionEject
# ---------------------------------------------------------------------------


class TestBudgetFractionEject:
    def test_no_eject_under_threshold(self):
        policy = BudgetFractionEject(
            check_after_pct=0.25,
            progress_fn=lambda ctx: False,  # no progress
        )
        d = policy.check(**_kwargs(_budget(0.01)))  # 4% spent
        assert not d.should_eject

    def test_eject_at_threshold_no_progress(self):
        policy = BudgetFractionEject(
            check_after_pct=0.20,
            progress_fn=lambda ctx: ctx.get("tool_calls", 0) > 0,
        )
        d = policy.check(**_kwargs(_budget(0.06), context={"tool_calls": 0}))  # 24%
        assert d.should_eject
        assert "no progress" in d.reason

    def test_no_eject_with_progress(self):
        policy = BudgetFractionEject(
            check_after_pct=0.20,
            progress_fn=lambda ctx: ctx.get("tool_calls", 0) > 0,
        )
        d = policy.check(**_kwargs(_budget(0.06), context={"tool_calls": 5}))
        assert not d.should_eject

    def test_fires_only_once(self):
        policy = BudgetFractionEject(
            check_after_pct=0.20,
            progress_fn=lambda ctx: False,
        )
        d1 = policy.check(**_kwargs(_budget(0.06)))
        assert d1.should_eject
        # Second check should not eject (already fired)
        d2 = policy.check(**_kwargs(_budget(0.10)))
        assert not d2.should_eject


# ---------------------------------------------------------------------------
# CompositeEject
# ---------------------------------------------------------------------------


class TestCompositeEject:
    def test_first_to_trigger_wins(self):
        policy = CompositeEject([
            BudgetExceededEject(),
            IdleTimeoutEject(timeout_s=9999),
        ])
        d = policy.check(**_kwargs(_budget(0.30)))
        assert d.should_eject
        assert "exhausted" in d.reason

    def test_no_eject_when_all_pass(self):
        policy = CompositeEject([
            BudgetExceededEject(),
            IdleTimeoutEject(timeout_s=9999),
        ])
        d = policy.check(**_kwargs(_budget(0.10), events=[{"type": "x"}]))
        assert not d.should_eject


class TestDefaultEjectPolicies:
    def test_returns_composite(self):
        p = default_eject_policies()
        assert isinstance(p, CompositeEject)
        assert len(p.policies) == 2
