from __future__ import annotations

import time

from tests.harness.test_runner import TestCase

SUITE = "test_timeouts"


def _response_within_timeout(ctx):
    start = time.monotonic()
    resp = ctx.client.prompt("Say hello in one sentence.", max_turns=2)
    elapsed = time.monotonic() - start
    assert isinstance(resp, dict)
    assert elapsed <= float(ctx.args.timeout), f"response exceeded timeout {ctx.args.timeout}s"
    return True, "response within timeout", {"elapsed": elapsed}


def _status_fast(ctx):
    start = time.monotonic()
    resp = ctx.client.status()
    elapsed = time.monotonic() - start
    assert resp.get("success") is True
    assert elapsed <= 2.0, f"status too slow: {elapsed}s"
    return True, "status fast", {"elapsed": elapsed}


def _context_fast(ctx):
    start = time.monotonic()
    resp = ctx.client.context_current()
    elapsed = time.monotonic() - start
    assert resp.get("success") is True
    assert elapsed <= 5.0, f"context too slow: {elapsed}s"
    return True, "context fast", {"elapsed": elapsed}


def get_tests():
    return [
        TestCase(SUITE, "response_within_timeout", "Prompt response within timeout", _response_within_timeout),
        TestCase(SUITE, "status_responds_fast", "Status endpoint latency", _status_fast),
        TestCase(SUITE, "context_responds_fast", "Context endpoint latency", _context_fast),
    ]
