"""Idle kernel reaping persists Python state and explains restoration to the model."""

import subprocess
import time
import unittest

from harness import Albedo, Provider, exclusive, python, text


def notes_for(app, session, cursor):
    return [
        event["text"]
        for event in app.stream_page(session, cursor)["events"]
        if event.get("type") == "note"
    ]


# Short enough that the test waits seconds, long enough to outlast a turn.
IDLE_SECONDS = 2


def latest_user(request):
    return [item["content"] for item in request["input"] if item.get("role") == "user"][
        -1
    ]


def kernels(daemon_pid):
    listing = subprocess.run(
        ["ps", "-o", "pid=,ppid=,command=", "-ax"], capture_output=True, text=True
    )
    processes = {}
    for line in listing.stdout.splitlines():
        fields = line.split(maxsplit=2)
        if len(fields) == 3 and fields[0].isdigit() and fields[1].isdigit():
            processes[int(fields[0])] = (int(fields[1]), fields[2])
    owned, frontier = [], [daemon_pid]
    while frontier:
        parent = frontier.pop()
        for pid, (ppid, command) in processes.items():
            if ppid == parent:
                frontier.append(pid)
                if "albedo_kernel.py" in command:
                    owned.append(pid)
    return owned


@exclusive
class IdleReapTests(unittest.TestCase):
    def setUp(self):
        def script(request):
            if request["input"][-1].get("role") == "user":
                prompt = latest_user(request)
                if prompt.startswith("remember"):
                    return python("answer = 7\nrows = [1, 2, 3]")
                if prompt.startswith("recall"):
                    return python("print('recalled', answer, rows)")
            return text("ok")

        provider = Provider(script)
        self.addCleanup(provider.close)
        self.provider = provider
        self.app = Albedo(
            provider,
            protocol="responses",
            prepare=lambda app: app.env.update(ALBEDO_IDLE_SECONDS=str(IDLE_SECONDS)),
        )
        self.app.__enter__()
        self.addCleanup(self.app.__exit__, None, None, None)

    def turn(self, session, prompt):
        self.app.prompt(session, prompt).close()
        self.app.idle(session)

    def test_idle_kernel_reaps_and_restores_variables_with_notices(self):
        app = self.app
        daemon = app.connection["pid"]
        existing = set(kernels(daemon))
        session = app.session()
        cursor = app.stream_page(session)
        self.assertEqual(
            set(kernels(daemon)) - existing,
            set(),
            "session must not eagerly open a kernel",
        )
        self.turn(session, "remember this")
        self.assertEqual(len(set(kernels(daemon)) - existing), 1)
        deadline = time.monotonic() + IDLE_SECONDS * 4
        while time.monotonic() < deadline and set(kernels(daemon)) - existing:
            time.sleep(0.1)
        self.assertEqual(
            set(kernels(daemon)) - existing, set(), "idle kernel outlived its limit"
        )
        notes = notes_for(app, session, cursor)
        self.assertIn("released", notes[-1])
        self.assertIn("2 variables saved to disk", notes[-1])
        self.assertTrue((app.home / "kernels" / f"{session}.state").exists())

        self.turn(session, "recall them")
        prompt = latest_user(self.provider.requests[-1]["request"])
        self.assertTrue(prompt.startswith("recall them<system-note>"), prompt)
        self.assertIn("restored from disk: answer, rows", prompt)
        self.assertTrue(
            any(
                "recalled 7 [1, 2, 3]" in item.get("output", "")
                for item in self.provider.requests[-1]["request"]["input"]
            )
        )
        notes = notes_for(app, session, cursor)
        self.assertTrue(
            any("restored 2 variables from disk" in note for note in notes), notes
        )