"""
Custom eject policy example.

Shows how to build a domain-specific eject policy by subclassing EjectPolicy.
The example ejects if a model has spent 30% of its budget but hasn't written
an output file yet.
"""

from pathlib import Path
from typing import Any

from agent_preflight import EjectDecision, EjectPolicy


class OutputFileEject(EjectPolicy):
    """Eject if no output file appears after spending check_pct of the budget."""

    def __init__(self, output_path: Path, check_pct: float = 0.30) -> None:
        self.output_path = output_path
        self.check_pct = check_pct
        self._fired = False

    def check(
        self,
        *,
        budget_status: Any,
        **_kwargs: Any,
    ) -> EjectDecision:
        if self._fired:
            return EjectDecision(should_eject=False)

        if budget_status.pct_spent >= self.check_pct:
            self._fired = True
            if not self.output_path.exists():
                return EjectDecision(
                    should_eject=True,
                    reason=(
                        f"no output at {self.output_path} after "
                        f"{budget_status.pct_spent*100:.0f}% budget spent"
                    ),
                )
        return EjectDecision(should_eject=False)


# Use it:
# eject = OutputFileEject(Path("/workspace/output/result.json"), check_pct=0.30)
# run_container(..., eject_policy=eject)
