"""Short-delay capability for agents that need a bounded wait between checks."""

from __future__ import annotations

import time
from typing import Any

from pydantic import BaseModel
from pydantic_ai.toolsets import FunctionToolset

from capabilities.base import CapabilityWithTools

MAX_SLEEP_SECONDS = 120


class SleepResult(BaseModel):
    """Return value of sleep_seconds."""

    slept_seconds: int


def sleep_seconds(seconds: int) -> SleepResult:
    """Sleep for a short bounded duration.

    Args->
        seconds: Number of seconds to sleep. Must be between 0 and 120 inclusive.

    Returns->
        {"slept_seconds": int} — access with result["slept_seconds"].
    """
    if seconds < 0:
        raise ValueError("seconds must be non-negative")
    if seconds > MAX_SLEEP_SECONDS:
        raise ValueError(f"seconds must be <= {MAX_SLEEP_SECONDS}")

    time.sleep(seconds)
    return SleepResult(slept_seconds=seconds)


class SleepCapability(CapabilityWithTools):
    def get_toolset(self) -> FunctionToolset[Any]:
        return FunctionToolset([sleep_seconds])

    def _instructions_header(self) -> str:
        return """\
## Wait Function

Call this as `await function_name(...)` inside `run_code`.

Use this function only for short bounded waits between checks. See the docstring for
the allowed range.\
"""


__all__ = ["MAX_SLEEP_SECONDS", "SleepCapability", "SleepResult", "sleep_seconds"]