"""Single-flight wake: N concurrent requests -> exactly one wake sequence."""

from __future__ import annotations

import asyncio

from config import DEPTH_AWAKE


async def test_ten_concurrent_requests_trigger_one_wake(pub, backend):
    backend.set_sleeping("text", True, level=2)
    backend.services["vllm-text"]["wake_delay"] = 0.15

    async def one(i: int):
        r = await pub.post("/v1/chat/completions", json={"model": "text", "seed": i})
        assert r.status_code == 200, r.text
        return r.json()

    results = await asyncio.gather(*(one(i) for i in range(10)))
    assert all(r["service"] == "vllm-text" for r in results)

    svc = backend.services["vllm-text"]
    assert svc["wake_up_calls"] == 1, backend.calls
    assert svc["reload_calls"] == 1
    assert svc["api_calls"] == 10


async def test_hold_timeout_does_not_cancel_the_wake(pub, backend, stack):
    """The first caller times out; the wake keeps going and finishes for the
    next caller (no second wake sequence, no half-awake proxy)."""
    backend.set_sleeping("text", True)
    backend.services["vllm-text"]["wake_delay"] = 0.3
    stack.cfg.hold_offload_s = 0.05

    r = await pub.post("/v1/chat/completions", json={"model": "text"})
    assert r.status_code == 503
    assert r.json()["error"]["sleep_depth"] == "offloaded"

    await asyncio.sleep(0.5)          # let the shielded wake task finish
    svc_state = stack.manager.services["text"]
    assert svc_state.depth == DEPTH_AWAKE
    assert backend.services["vllm-text"]["wake_up_calls"] == 1

    r = await pub.post("/v1/chat/completions", json={"model": "text"})
    assert r.status_code == 200
    assert backend.services["vllm-text"]["wake_up_calls"] == 1   # cached awake


async def test_two_services_wake_concurrently_without_cross_talk(pub, backend):
    backend.set_sleeping("text", True)
    backend.set_sleeping("embed", True)
    backend.services["vllm-text"]["wake_delay"] = 0.2
    backend.services["vllm-embed"]["wake_delay"] = 0.05

    chat = pub.post("/v1/chat/completions", json={"model": "text"})
    embed = pub.post("/v1/embeddings", json={"input": "hi"})
    r1, r2 = await asyncio.gather(chat, embed)
    assert r1.status_code == 200 and r1.json()["service"] == "vllm-text"
    assert r2.status_code == 200 and r2.json()["service"] == "vllm-embed"
    assert backend.services["vllm-text"]["wake_up_calls"] == 1
    assert backend.services["vllm-embed"]["wake_up_calls"] == 1


async def test_wake_state_is_cached_between_requests(pub, backend, stack):
    stack.cfg.state_cache_ttl = 30.0
    backend.set_sleeping("ocr", True)
    assert (await pub.post("/v1/chat/completions", json={"model": "ocr"})).status_code == 200
    probes_after_first = backend.count("GET vllm-ocr/is_sleeping")
    assert probes_after_first >= 1
    for _ in range(5):
        assert (await pub.post("/v1/chat/completions", json={"model": "ocr"})).status_code == 200
    # no extra probe per request while the awake state is cached
    assert backend.count("GET vllm-ocr/is_sleeping") == probes_after_first


async def test_expired_cache_reprobes(pub, backend, stack):
    stack.cfg.state_cache_ttl = 0.0
    for _ in range(3):
        assert (await pub.post("/v1/chat/completions", json={"model": "ocr"})).status_code == 200
    assert backend.count("GET vllm-ocr/is_sleeping") >= 3


async def test_admin_and_request_path_share_one_lock(adm, pub, backend):
    """A wake driven from the admin port is joined by the public request path."""
    backend.set_sleeping("text", True)
    backend.services["vllm-text"]["wake_delay"] = 0.2

    admin_wake = asyncio.create_task(adm.post("/admin/wake/text"))
    await asyncio.sleep(0.05)
    r = await pub.post("/v1/chat/completions", json={"model": "text"})
    admin_result = await admin_wake

    assert r.status_code == 200
    assert admin_result.status_code == 200
    assert backend.services["vllm-text"]["wake_up_calls"] == 1