"""Router-restart-mid-wake recovery (plan 6.4, E2E case 12).

`is_sleeping=false` does NOT mean "reload done": if the previous router died
between POST /wake_up and collective_rpc reload_weights, the backend happily
reports awake+healthy while serving garbage.  The persisted wake-intent file
is what lets the new router tell the two apart.
"""

from __future__ import annotations

import asyncio
import json
import os

import httpx

from app import build_apps
from config import DEPTH_OFFLOADED, clone, load_config
from conftest import raw_asgi
from services import ServiceManager


def write_state(path, **services) -> None:
    payload = {"version": 1, "updated": 0.0, "services": {
        key: {"depth": value.get("depth"),
              "wake_in_progress": value.get("wake_in_progress", False),
              "level": value.get("level", 0),
              "pending_reload": value.get("pending_reload", False)}
        for key, value in services.items()
    }}
    with open(path, "w", encoding="utf-8") as fh:
        json.dump(payload, fh)


def read_state(path) -> dict:
    with open(path, encoding="utf-8") as fh:
        return json.load(fh)["services"]


async def _make(cfg, backend, state_file):
    public_app, admin_app, manager = build_apps(cfg, transport=backend)
    try:
        yield public_app, admin_app, manager
    finally:
        await manager.proxy_client.aclose()
        await manager.ctrl_client.aclose()


async def test_startup_completes_interrupted_wake(tmp_path, backend):
    """The E2E failure: old router died after wake_up, before reload_weights.
    The new router must not proxy until the sequence has been completed."""
    state_file = str(tmp_path / "state.json")
    write_state(state_file, text={"depth": "offloaded", "wake_in_progress": True,
                                  "level": 2})
    # backend *looks* perfectly healthy and awake -- that is the lie
    backend.services["vllm-text"]["sleeping"] = False

    cfg = clone(load_config())
    cfg.state_file = state_file
    cfg.idle_enabled = False
    cfg.state_cache_ttl = 0.0
    async for _public, _admin, manager in _make(cfg, backend, state_file):
        manager.start()
        await asyncio.sleep(0.3)          # let the recovery task run

        text = manager.services["text"]
        assert text.pending_reload is False
        assert text.depth == "awake"
        svc = backend.services["vllm-text"]
        # the interrupted sequence was completed, not skipped
        assert svc["reload_calls"] == 1
        assert svc["reset_calls"] == 1
        # and the file no longer claims a wake in flight
        assert read_state(state_file)["text"]["wake_in_progress"] is False


async def test_interrupted_wake_never_proxies_before_recovery(tmp_path, backend, pub):
    """A request arriving in the recovery window must not be proxied on the
    strength of is_sleeping=false alone."""
    state_file = str(tmp_path / "state.json")
    write_state(state_file, ocr={"depth": "offloaded", "wake_in_progress": True, "level": 2})
    backend.services["vllm-ocr"]["sleeping"] = False     # awake-looking
    backend.services["vllm-ocr"]["api_delay"] = 0.0
    backend.calls.clear()

    stack_cfg = pub  # unused; keeps the fixture ordering simple
    del stack_cfg
    cfg = clone(load_config())
    cfg.state_file = state_file
    cfg.idle_enabled = False
    async for public_app, _admin, manager in _make(cfg, backend, state_file):
        manager.start()
        # fire the request immediately, before recovery has finished probing
        task = asyncio.create_task(_post(public_app, "/v1/chat/completions",
                                         {"model": "ocr"}))
        await asyncio.sleep(0.05)
        ocr = backend.services["vllm-ocr"]
        assert ocr["reload_calls"] >= 0
        response = await task
        assert response.status_code == 200
        # proxied only after a completed reload sequence
        assert backend.last("POST vllm-ocr/v1/chat/completions") > \
            backend.last("POST vllm-ocr/reset_prefix_cache")
        assert manager.services["ocr"].pending_reload is False


async def _post(app, path, payload) -> httpx.Response:
    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app),
                                 base_url="http://router.test") as client:
        return await client.post(path, json=payload)


async def test_startup_clears_state_when_unreachable(tmp_path, backend):
    """wake_in_progress + unreachable backend = fresh boot (weights are fresh);
    nothing to complete."""
    state_file = str(tmp_path / "state.json")
    write_state(state_file, text={"depth": "offloaded", "wake_in_progress": True, "level": 2})
    backend.services["vllm-text"]["reachable"] = False

    cfg = clone(load_config())
    cfg.state_file = state_file
    cfg.idle_enabled = False
    async for _public, _admin, manager in _make(cfg, backend, state_file):
        manager.start()
        await asyncio.sleep(0.2)
        text = manager.services["text"]
        assert text.pending_reload is False
        assert backend.services["vllm-text"]["reload_calls"] == 0
        assert read_state(state_file)["text"]["wake_in_progress"] is False


async def test_startup_clears_state_when_still_sleeping(tmp_path, backend):
    """wake_in_progress but the backend never woke: the normal request path
    will run a full wake when traffic arrives."""
    state_file = str(tmp_path / "state.json")
    write_state(state_file, embed={"depth": "sleeping", "wake_in_progress": True, "level": 1})
    backend.set_sleeping("embed", True, level=1)

    cfg = clone(load_config())
    cfg.state_file = state_file
    cfg.idle_enabled = False
    async for _public, _admin, manager in _make(cfg, backend, state_file):
        manager.start()
        await asyncio.sleep(0.2)
        embed = manager.services["embed"]
        assert embed.pending_reload is False
        assert embed.depth == "sleeping"          # restored from the file
        assert backend.services["vllm-embed"]["reload_calls"] == 0


async def test_no_reload_when_no_wake_in_progress(tmp_path, backend, pub):
    """Plain restart with a settled backend: no reload, no state churn."""
    state_file = str(tmp_path / "state.json")
    write_state(state_file,
                text={"depth": "awake", "wake_in_progress": False},
                ocr={"depth": "offloaded", "wake_in_progress": False, "level": 2})

    cfg = clone(load_config())
    cfg.state_file = state_file
    cfg.idle_enabled = False
    async for _public, _admin, manager in _make(cfg, backend, state_file):
        manager.start()
        await asyncio.sleep(0.1)
        assert manager.services["ocr"].depth == "offloaded"
        assert backend.count("POST vllm-ocr/wake_up") == 0
        assert backend.count("POST vllm-text/collective_rpc") == 0

        # the awake service serves immediately, straight through the proxy
        status, body = await raw_asgi(_public_app(manager), "POST",
                                      "/v1/chat/completions",
                                      headers=[("content-type", "application/json")],
                                      body=b'{"model": "text"}')
        assert status == 200, body
        svc = backend.services["vllm-text"]
        assert svc["wake_up_calls"] == 0          # no pointless wake
        assert svc["reload_calls"] == 0           # no pointless reload


def _public_app(manager):
    from routing import build_public_app
    return build_public_app(manager)


async def test_interrupted_wake_end_to_end(tmp_path, backend):
    """In-process replay of E2E case 12: the wake is killed after wake_up has
    flipped is_sleeping but before reload_weights; a new manager on the same
    state file must complete the sequence before proxying."""
    state_file = str(tmp_path / "state.json")
    cfg = clone(load_config())
    cfg.state_file = state_file
    cfg.idle_enabled = False
    cfg.wake_health_poll_s = 0.005

    backend.set_sleeping("text", True, level=2)
    backend.services["vllm-text"]["wake_delay"] = 0.4
    backend.services["vllm-text"]["wake_flip_midway"] = True

    first = ServiceManager(cfg, transport=backend)
    try:
        task = asyncio.create_task(first.ensure_awake("text"))
        await asyncio.sleep(0.3)              # is_sleeping flipped at t+0.2s
        task.cancel()                          # ... and the router "dies"
        await asyncio.gather(task, return_exceptions=True)
    finally:
        await first.proxy_client.aclose()
        await first.ctrl_client.aclose()

    assert backend.services["vllm-text"]["sleeping"] is False   # looks awake
    assert backend.services["vllm-text"]["reload_calls"] == 0   # ... but isn't
    assert read_state(state_file)["text"]["wake_in_progress"] is True

    backend.calls.clear()
    async for public_app, _admin, manager in _make(cfg, backend, state_file):
        manager.start()
        response = await _post(public_app, "/v1/chat/completions", {"model": "text"})
        assert response.status_code == 200
        svc = backend.services["vllm-text"]
        assert svc["reload_calls"] == 1        # sequence was completed ...
        assert backend.last("POST vllm-text/v1/chat/completions") > \
            backend.last("POST vllm-text/reset_prefix_cache")   # ... first
        assert manager.services["text"].pending_reload is False


async def test_wake_intent_is_written_and_cleared(tmp_path, backend):
    """The file says wake_in_progress=true while a wake is running and false
    once it completes -- which is exactly what a restart in the middle leaves
    behind."""
    state_file = str(tmp_path / "state.json")
    cfg = clone(load_config())
    cfg.state_file = state_file
    cfg.idle_enabled = False
    cfg.wake_health_poll_s = 0.005
    backend.set_sleeping("text", True, level=2)
    backend.services["vllm-text"]["wake_delay"] = 0.2

    manager = ServiceManager(cfg, transport=backend)
    try:
        task = asyncio.create_task(manager.ensure_awake("text"))
        await asyncio.sleep(0.05)                 # mid-wake
        assert read_state(state_file)["text"]["wake_in_progress"] is True
        outcome = await task
        assert outcome.ok
        entry = read_state(state_file)["text"]
        assert entry["wake_in_progress"] is False
        assert entry["depth"] == "awake"
    finally:
        await manager.proxy_client.aclose()
        await manager.ctrl_client.aclose()


async def test_failed_wake_clears_the_intent(tmp_path, backend):
    state_file = str(tmp_path / "state.json")
    cfg = clone(load_config())
    cfg.state_file = state_file
    cfg.idle_enabled = False
    cfg.wake_health_poll_s = 0.005
    backend.set_sleeping("text", True)
    backend.services["vllm-text"]["wake_fails"] = 9

    manager = ServiceManager(cfg, transport=backend)
    try:
        outcome = await manager.ensure_awake("text")
        assert not outcome.ok
        entry = read_state(state_file)["text"]
        assert entry["wake_in_progress"] is False
    finally:
        await manager.proxy_client.aclose()
        await manager.ctrl_client.aclose()


async def test_state_file_failure_is_not_fatal(tmp_path, backend):
    """An unwritable state path disables persistence, never serving."""
    cfg = clone(load_config())
    cfg.state_file = str(tmp_path / "no-such-dir" / "state.json")
    cfg.idle_enabled = False
    async for public_app, _admin, manager in _make(cfg, backend, cfg.state_file):
        manager.start()
        response = await _post(public_app, "/v1/chat/completions", {"model": "text"})
        assert response.status_code == 200
        assert not os.path.exists(cfg.state_file)


async def test_admin_reports_recovery(tmp_path, backend):
    state_file = str(tmp_path / "state.json")
    write_state(state_file, ocr={"depth": "offloaded", "wake_in_progress": True, "level": 2})
    backend.services["vllm-ocr"]["sleeping"] = False

    cfg = clone(load_config())
    cfg.state_file = state_file
    cfg.idle_enabled = False
    async for _public, admin_app, manager in _make(cfg, backend, state_file):
        async with httpx.AsyncClient(transport=httpx.ASGITransport(app=admin_app),
                                     base_url="http://router.test") as client:
            r = await client.get("/admin/status")
            assert r.status_code == 200
            assert r.json()["services"]["ocr"]["wake_recovery_pending"] is True