"""Native Codex quota refresh: one initial read, non-blocking notifications."""
import asyncio

import pytest

from device_gateway.control.codex import CodexAdapter
from device_gateway.control.codex_quota import INVALID
from device_gateway.control.models import Project
from device_gateway.control.projects import ProjectRegistry


def snapshot(used, resets=1700000000, minutes=300):
    return {"rateLimitsByLimitId": {"codex": {"primary": {
        "usedPercent": used, "resetsAt": resets, "windowDurationMins": minutes}}}}


class Native:
    """Records account-level reads; can block to prove reader independence."""

    generation = 1

    def __init__(self):
        self.calls = []
        self.quota_reads = 0
        self.read_started = None
        self.read_release = None
        self.payload = snapshot(28)

    async def start(self):
        pass

    async def write(self, value):
        pass

    async def close(self):
        pass

    async def call(self, method, params):
        self.calls.append((method, params))
        if method == "initialize":
            return {}
        assert method == "account/rateLimits/read"
        self.quota_reads += 1
        if self.read_started is not None:
            self.read_started.set()
            await self.read_release.wait()
        return self.payload

    def initialization_reads(self):
        return sum(method == "account/rateLimits/read" for method, _ in self.calls)


def adapter(tmp_path):
    projects = ProjectRegistry(tmp_path / "projects.json")
    projects.add(Project(id="p", name="P", path=str(tmp_path)))
    native = Native()
    return CodexAdapter(projects, native), native


def run(coro):
    return asyncio.run(coro)


def test_initial_read_once_per_successful_initialize(tmp_path):
    instance, native = adapter(tmp_path)

    async def scenario():
        await instance.ready()
        assert native.initialization_reads() == 1
        assert instance.quota.valid and instance.quota.remaining_percent == 72
        # A second ready() in the same generation must not re-read.
        await instance.ready()
        assert native.initialization_reads() == 1

    run(scenario())


def test_updated_notification_schedules_one_refresh(tmp_path):
    instance, native = adapter(tmp_path)

    async def scenario():
        await instance.ready()
        assert native.quota_reads == 1
        await instance.event({"method": "account/rateLimits/updated",
                              "params": {"rateLimits": {"primary": {"usedPercent": 40,
                                                                    "resetsAt": 1700000000,
                                                                    "windowDurationMins": 300}}}})
        # Sparse values apply immediately, before the refetch lands.
        assert instance.quota.remaining_percent == 60
        if instance.quota_refresh_task:
            await instance.quota_refresh_task
        assert native.quota_reads == 2

    run(scenario())


def test_notification_does_not_block_socket_reader(tmp_path):
    instance, native = adapter(tmp_path)

    async def scenario():
        await instance.ready()
        native.read_started = asyncio.Event()
        native.read_release = asyncio.Event()
        native.payload = snapshot(10)
        # The handler must return without awaiting the RPC response.
        await asyncio.wait_for(instance.event({"method": "account/rateLimits/updated",
            "params": {"rateLimits": {"primary": {"usedPercent": 10}}}}), timeout=1)
        assert native.read_started is not None
        # Reader is free: a second unrelated notification is processed too.
        await asyncio.wait_for(instance.event({"method": "account/rateLimits/updated",
            "params": {"rateLimits": {"primary": {"usedPercent": 20}}}}), timeout=1)
        native.read_release.set()
        if instance.quota_refresh_task:
            await instance.quota_refresh_task

    run(scenario())


def test_notification_storm_creates_bounded_concurrent_reads(tmp_path):
    instance, native = adapter(tmp_path)

    async def scenario():
        await instance.ready()
        baseline = native.quota_reads
        native.read_started = asyncio.Event()
        native.read_release = asyncio.Event()
        for index in range(25):
            await instance.event({"method": "account/rateLimits/updated",
                                  "params": {"rateLimits": {"primary": {"usedPercent": index}}}})
        # At most one in-flight refresh exists regardless of notification count.
        assert sum(1 for task in [instance.quota_refresh_task] if task and not task.done()) <= 1
        native.read_started.set()
        native.read_release.set()
        if instance.quota_refresh_task:
            await instance.quota_refresh_task
        await asyncio.sleep(0)
        assert native.quota_reads - baseline <= 1

    run(scenario())


def test_refresh_failure_keeps_last_known_value(tmp_path):
    instance, native = adapter(tmp_path)

    async def scenario():
        await instance.ready()
        assert instance.quota.remaining_percent == 72

        async def failing(method, params):
            raise RuntimeError("provider unavailable")

        native.call = failing
        await instance.refresh_quota()
        # Quota is preserved rather than cleared; unavailable is reported only
        # when nothing was ever observed.
        assert instance.quota.valid and instance.quota.remaining_percent == 72

    run(scenario())


def test_disconnect_clears_quota_and_cancels_refresh(tmp_path):
    instance, native = adapter(tmp_path)

    async def scenario():
        await instance.ready()
        instance.quota_refresh_task = asyncio.create_task(asyncio.sleep(30))
        instance.disconnected()
        assert instance.quota == INVALID
        assert instance.quota_refresh_task is None

    run(scenario())


def test_initialize_absent_quota_is_invalid_not_fatal(tmp_path):
    instance, native = adapter(tmp_path)
    native.payload = {}

    async def scenario():
        await instance.ready()
        assert instance.quota == INVALID
        assert instance.quota.wire() == {"codex_quota_valid": False,
            "codex_remaining_percent": 0, "codex_resets_at": None, "codex_window_minutes": None}

    run(scenario())


# ---------------------------------------------------------------------------
# Periodic refresh: the cache must catch up even without a native event.
# ---------------------------------------------------------------------------

def test_periodic_refresh_updates_stale_cache(tmp_path):
    """51 -> provider moves to 43 -> the next server_snapshot reports 43."""
    instance, native = adapter(tmp_path)
    native.payload = snapshot(49)  # remaining 51

    async def scenario():
        await instance.ready()
        assert instance.quota.remaining_percent == 51
        # Simulate a tiny cadence so the loop fires quickly in-test.
        instance.quota_periodic_seconds = 0.02
        native.payload = snapshot(57)  # provider later reports 43
        await asyncio.sleep(0.06)
        instance._stop_quota_periodic()
        assert instance.quota.remaining_percent == 43
        # The device-facing snapshot reads only the cache.
        from device_gateway.app import device_snapshot_payload
        assert device_snapshot_payload.__name__  # import sanity

    run(scenario())


def test_periodic_refresh_is_not_per_snapshot(tmp_path):
    """A long cadence means many snapshots produce no extra native RPCs."""
    instance, native = adapter(tmp_path)

    async def scenario():
        await instance.ready()
        baseline = native.quota_reads
        instance.quota_periodic_seconds = 30
        # Repeated cache reads (what a 2 s snapshot cadence does) never touch RPC.
        for _ in range(15):
            _ = instance.quota
        await asyncio.sleep(0.05)
        assert native.quota_reads == baseline
        instance._stop_quota_periodic()

    run(scenario())


def test_periodic_refresh_coalesces_with_event_refresh(tmp_path):
    """At most one in-flight quota read exists even with a timer and events."""
    instance, native = adapter(tmp_path)

    async def scenario():
        await instance.ready()
        native.read_started = asyncio.Event()
        native.read_release = asyncio.Event()
        native.read_started.clear()
        instance.quota_periodic_seconds = 0.01
        await instance.event({"method": "account/rateLimits/updated",
                              "params": {"rateLimits": {"primary": {"usedPercent": 30}}}})
        in_flight_before = native.quota_reads
        await asyncio.sleep(0.05)
        # The lock serializes every refresh; the event refresh and the timer
        # refresh cannot both add reads beyond the serialized sequence.
        assert native.quota_reads - in_flight_before <= 2
        native.read_started.set()
        native.read_release.set()
        instance._stop_quota_periodic()
        if instance.quota_refresh_task:
            await instance.quota_refresh_task

    run(scenario())


def test_quota_updated_at_records_success_only(tmp_path):
    instance, native = adapter(tmp_path)

    async def scenario():
        assert instance.quota_updated_at is None
        await instance.ready()
        first = instance.quota_updated_at
        assert first is not None

        async def failing(method, params):
            raise RuntimeError("provider unavailable")

        native.call = failing
        await instance.refresh_quota()
        # A failed read never advances the timestamp; the cache is preserved.
        assert instance.quota_updated_at == first
        assert instance.quota.remaining_percent == 72
        instance._stop_quota_periodic()

    run(scenario())


def test_periodic_refresh_failure_keeps_agent_alive(tmp_path):
    instance, native = adapter(tmp_path)

    async def scenario():
        await instance.ready()
        held = instance.quota.remaining_percent

        async def failing(method, params):
            if method == "account/rateLimits/read":
                raise RuntimeError("provider unavailable")
            return {}

        native.call = failing
        instance.quota_periodic_seconds = 0.01
        await asyncio.sleep(0.05)
        instance._stop_quota_periodic()
        # Repeated failures neither clear the cache nor raise out of the loop.
        assert instance.quota.valid and instance.quota.remaining_percent == held
        # Agent calls still work after the failures.
        await instance.call("thread/list", {})

    run(scenario())


def test_periodic_refresh_clean_shutdown(tmp_path):
    instance, native = adapter(tmp_path)

    async def scenario():
        await instance.ready()
        task = instance.quota_periodic_task
        assert task is not None and not task.done()
        instance.disconnected()
        assert instance.quota_periodic_task is None
        await asyncio.sleep(0)  # let the cancellation settle
        assert task.cancelled() or task.done()
        assert instance.quota == INVALID

    run(scenario())


def test_periodic_refresh_disabled_when_interval_zero(tmp_path):
    instance, native = adapter(tmp_path)
    instance.quota_periodic_seconds = 0

    async def scenario():
        await instance.ready()
        assert instance.quota_periodic_task is None
        assert native.quota_reads == 1  # only the initialize read

    run(scenario())