"""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())

def test_default_cadence_and_stale_display(tmp_path, monkeypatch):
    monkeypatch.delenv('GATEWAY_CODEX_QUOTA_REFRESH_SECONDS', raising=False)
    monkeypatch.delenv('GATEWAY_CODEX_QUOTA_TIMEOUT_SECONDS', raising=False)
    instance, native = adapter(tmp_path)
    assert instance.quota_periodic_seconds == 15
    # The read reaches OpenAI through the local daemon and is routinely slower
    # than a few seconds, so the default budget must not be a short cap.
    assert instance.quota_timeout_seconds == 45

    async def scenario():
        await instance.ready()
        assert instance.quota_for_display().valid
        # A reading stays current across a full attempt plus a retry budget.
        instance.quota_updated_at -= instance.quota_stale_after() + 1
        assert instance.quota_for_display() == INVALID
        await instance.refresh_quota()
        assert instance.quota_for_display().valid
        instance._stop_quota_periodic()

    run(scenario())


def test_stale_threshold_covers_a_slow_read(tmp_path, monkeypatch):
    """A read that takes almost the whole timeout must not blank the device.

    Regression: a hardcoded 45 s staleness threshold against a 5 s timeout
    expired the cache mid-flight once the provider slowed to >5 s per read.
    """
    monkeypatch.delenv('GATEWAY_CODEX_QUOTA_REFRESH_SECONDS', raising=False)
    monkeypatch.delenv('GATEWAY_CODEX_QUOTA_TIMEOUT_SECONDS', raising=False)
    instance, native = adapter(tmp_path)

    async def scenario():
        await instance.ready()
        assert instance.quota_for_display().valid
        # The budget must exceed one read timeout: a reading taken just before a
        # slow attempt must survive until that attempt can land.
        assert instance.quota_stale_after() > instance.quota_timeout_seconds
        instance.quota_updated_at -= instance.quota_timeout_seconds + 1
        assert instance.quota_for_display().valid
        instance._stop_quota_periodic()

    run(scenario())


def test_slow_read_within_timeout_keeps_quota(tmp_path, monkeypatch):
    """A read slower than the old 5 s cap still lands and stays valid."""
    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()

        # A 45 s budget must outlast a read that the old 5 s cap would abort.
        async def bounded(method, params):
            if method == "initialize":
                return {}
            native.quota_reads += 1
            native.read_started.set()
            await asyncio.wait_for(native.read_release.wait(), 1)
            return snapshot(10)

        native.call = bounded
        task = asyncio.create_task(instance.refresh_quota())
        await native.read_started.wait()
        # Mid-flight the previous reading is still the displayed one.
        assert instance.quota_for_display().valid
        native.read_release.set()
        await task
        assert native.quota_reads == baseline + 1
        assert instance.quota.valid and instance.quota.remaining_percent == 90
        assert instance.quota_for_display().valid
        instance._stop_quota_periodic()

    run(scenario())


def test_failed_reads_back_off_and_reset_on_success(tmp_path):
    instance, native = adapter(tmp_path)
    instance.quota_backoff_seconds = 15
    instance.quota_backoff_max_seconds = 120

    async def scenario():
        await instance.ready()
        assert instance.quota_backoff_delay() == 0.0
        healthy = native.call

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

        native.call = failing
        expected = [15, 30, 60, 120, 120]
        for delay in expected:
            await instance.refresh_quota()
            assert instance.quota_backoff_delay() == delay
        # The cache survives the whole failure streak.
        assert instance.quota.valid and instance.quota.remaining_percent == 72

        native.call = healthy
        await instance.refresh_quota()
        assert instance.quota_failures == 0
        assert instance.quota_backoff_delay() == 0.0
        instance._stop_quota_periodic()

    run(scenario())


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

    async def scenario():
        await instance.ready()

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

        native.call = failing
        await instance.refresh_quota()
        assert instance.quota_failures == 1
        await instance.event({"method": "account/rateLimits/updated",
                              "params": {"rateLimits": {"primary": {"usedPercent": 40,
                                                                    "resetsAt": 1700000000,
                                                                    "windowDurationMins": 300}}}})
        assert instance.quota_failures == 0
        assert instance.quota_for_display().valid
        if instance.quota_refresh_task:
            instance.quota_refresh_task.cancel()
        instance._stop_quota_periodic()

    run(scenario())

def test_late_query_cannot_overwrite_newer_notification(tmp_path):
    instance,native=adapter(tmp_path)
    async def scenario():
        await instance.ready()
        native.read_started=asyncio.Event();native.read_release=asyncio.Event()
        query=asyncio.create_task(instance.refresh_quota())
        await native.read_started.wait()
        await instance.event({'method':'account/rateLimits/updated','params':{'rateLimits':{'primary':{'usedPercent':80}}}})
        native.read_release.set();await query
        assert instance.quota.remaining_percent==20
        instance._stop_quota_periodic()
    run(scenario())