import asyncio
import time

from langchain_quickjs import CodeInterpreterMiddleware


# Minimal stand-ins for the deepagents ModelRequest the hooks expect.
class _Req:
    tools: list = []
    system_message = None
    def override(self, **kw):
        return self


async def main() -> None:
    # ptc=["eval"] makes awrap_model_call -> _prepare_for_call build the slot
    # via registry.get() on the cache miss. Without ptc the hook early-returns.
    mw = CodeInterpreterMiddleware(ptc=["eval"])
    req = _Req()

    async def handler(request):
        return "ok"

    async def heartbeat() -> None:
        last = time.monotonic()
        while True:
            await asyncio.sleep(0.01)
            gap = time.monotonic() - last
            last = time.monotonic()
            tag = "   <-- loop was blocked" if gap > 0.05 else ""
            print(f"heartbeat gap {gap * 1000:6.1f} ms{tag}")

    hb = asyncio.create_task(heartbeat())
    await asyncio.sleep(0.15)  # healthy baseline

    # 1) Build path: awrap_model_call -> _prepare_for_call -> registry.get()
    #    -> _build_slot_locked -> worker.run_sync()  (blocks on the loop thread)
    print(">>> await mw.awrap_model_call(req, handler)  (build: worker.run_sync)")
    await mw.awrap_model_call(req, handler)
    await asyncio.sleep(0.15)

    # Prime the worker with allocs so close()'s gc.collect() has real work
    # (an idle worker thread joins near-instantly; GC under pressure is the
    # production condition).
    thread_id = mw._fallback_thread_id
    mw._registry.get(thread_id).eval_sync(
        "let a = []; for (let i = 0; i < 2_000_000; i++) a.push({x: i}); a.length = 0;"
    )

    # 2) Evict path: aafter_agent -> registry.aevict() -> _aclose_slot
    #    -> worker.close() (thread.join + gc.collect)  (blocks on the loop thread)
    print(">>> await mw.aafter_agent(state, runtime)  (evict: worker.close)")
    state = {"_quickjs_snapshot_payload": None}  # thread mode w/o payload -> evict branch
    await mw.aafter_agent(state, object())  # runtime is unused by the hook
    print(">>> aafter_agent returned")

    hb.cancel()


asyncio.run(main())

def close(self) -> None:
    """Stop the worker thread and join. Idempotent."""
    ...
    self._loop.call_soon_threadsafe(self._loop.stop)
    self._thread.join(timeout=5.0)   # ← blocks the calling (event-loop) thread
    self._loop = None
    self._thread = None

async def _aclose_slot(self, slot: _Slot) -> None:
    with contextlib.suppress(Exception):
        await slot.repl.aclose()
    with contextlib.suppress(Exception):
        await slot.worker.run_async(_aclose_runtime(slot.runtime))
    slot.worker.close()   # ← blocking, on the loop thread

def _build_slot_locked(self, thread_id: str) -> _Slot:
    ...
    worker = ThreadWorker(name=name)
    runtime = worker.run_sync(self._acreate_runtime())  # ← blocking init, on the loop thread
    repl = _ThreadREPL(worker, runtime, ...)
    return _Slot(worker=worker, runtime=runtime, repl=repl)

# abefore_agent:379
repl = self._registry.get(thread_id)          # blocking on miss
# _prepare_for_call:462 (called from awrap_model_call:412)
repl = self._registry.get(thread_id)          # blocking on miss