"""Memory profiling tests for the smesh WASM web app.

Requires an instrumented build:
  MOXIEROOT=../moxie moxie build -target wasm -print-allocs "web/wasm/app" \
    -o web/static/app.wasm ./web/wasm/app/

Run:
  pytest test/test_memory_profile.py --headed
"""

import time
import json
import pytest


def _trigger_stats(h):
    h.js("window.__moxie_trigger_mem_stats && window.__moxie_trigger_mem_stats()")
    time.sleep(0.3)


def _mem(h):
    _trigger_stats(h)
    return h.js("return window.__moxie_mem_stats ? window.__moxie_mem_stats() : null") or {}


def _counters(h):
    _trigger_stats(h)
    return h.js("return window.__moxie_read_alloc_counters ? window.__moxie_read_alloc_counters() : null") or {}


@pytest.fixture(scope="module")
def app(browser, relay, h):
    """Navigate to the app and wait until WASM is ready."""
    browser.get(relay["url"])
    h.wait_wasm_ready(timeout=20)
    # Allow initial render to settle.
    time.sleep(2)
    return browser


@pytest.mark.memory
class TestMemoryProfile:

    def test_mem_stats_available(self, app, h):
        """Verify instrumented WASM exports are accessible."""
        stats = _mem(h)
        assert stats is not None, "mem_stats returned None — is this an instrumented build?"
        assert "totalAlloc" in stats or "heapPtr" in stats, f"unexpected stats shape: {stats}"

    def test_baseline(self, app, h):
        """Record initial heap state."""
        stats = _mem(h)
        heap_mb = stats.get("heapPtr", 0) / (1024 * 1024)
        total_mb = stats.get("totalAlloc", 0) / (1024 * 1024)
        print(f"\n[baseline] heapPtr={heap_mb:.1f}MB totalAlloc={total_mb:.1f}MB mallocs={stats.get('mallocs', 0)}")
        # Baseline should be well under 64MB.
        assert heap_mb < 64, f"baseline heap {heap_mb:.1f}MB exceeds 64MB"

    def test_dom_handles_bounded(self, app, h):
        """DOM handle count must stay under 3000."""
        count = h.js("return window.__moxie_alloc_log ? 'logging' : 'no_log'")
        # If no instrumentation log, skip quantitative check.
        if count == "no_log":
            pytest.skip("__moxie_alloc_log not initialised — set globalThis.__moxie_instrument=true")

        # Enable instrumentation and wait for threshold events.
        h.js("globalThis.__moxie_instrument = true;")
        time.sleep(1)
        log = h.js("return (globalThis.__moxie_alloc_log || []).slice(-10)")
        if log:
            max_dom = max((e.get("value", 0) for e in log if e.get("source") == "dom"), default=0)
            assert max_dom < 3000, f"DOM handle count reached {max_dom} (limit 3000)"

    def test_idle_alloc_rate(self, app, h):
        """After subscriptions are closed, allocation rate must be < 1KB/s.

        This tests that the rotate-map eviction is working: a live session
        with no incoming events should allocate nearly nothing.
        """
        before = _mem(h)
        if not before.get("totalAlloc"):
            pytest.skip("totalAlloc not available in this build")

        # Close the feed subscription to achieve idle state.
        h.js("window.__moxie_trigger_mem_stats && window.__moxie_trigger_mem_stats()")
        time.sleep(30)

        after = _mem(h)
        delta_bytes = after.get("totalAlloc", 0) - before.get("totalAlloc", 0)
        rate_bytes_per_sec = delta_bytes / 30

        print(f"\n[idle_rate] delta={delta_bytes}B over 30s = {rate_bytes_per_sec:.0f}B/s")
        assert rate_bytes_per_sec < 1024, (
            f"idle alloc rate {rate_bytes_per_sec:.0f}B/s exceeds 1KB/s — "
            f"rotate-map may not be active or relay not quiesced"
        )

    def test_alloc_counters_present(self, app, h):
        """Verify per-site counters are populated when -print-allocs was used."""
        counters = _counters(h)
        if not counters or counters.get("n", 0) == 0:
            pytest.skip("Alloc counters not available — build without -print-allocs?")

        n = counters["n"]
        counts = counters.get("counts", [])
        sizes = counters.get("sizes", [])

        assert n > 0, "n should be positive"
        assert len(counts) >= n
        assert len(sizes) >= n

        total_allocs = sum(counts[:n])
        total_bytes = sum(sizes[:n])
        print(f"\n[counters] {n} sites, {total_allocs} allocs, {total_bytes}B total")
        # At least some allocs should have happened.
        assert total_allocs > 0, "no allocations recorded"

    def test_heap_under_restart_threshold(self, app, h):
        """After normal use heapPtr should be well under the 128MB restart threshold."""
        stats = _mem(h)
        heap_bytes = stats.get("heapPtr", 0)
        heap_mb = heap_bytes / (1024 * 1024)
        print(f"\n[heap_check] heapPtr={heap_mb:.1f}MB")
        assert heap_mb < 64, (
            f"heap {heap_mb:.1f}MB exceeds 64MB after normal use — "
            f"rotate-map may not be reducing growth"
        )