"""A session's kernel outlives its daemon and its bridge.

The kernel runs detached behind a bridge process; the daemon reattaches to it
after a restart (graceful or crash) and after the bridge dies mid-cell, so the
namespace, the kernel process, and an in-flight cell's result all survive.
What the kernel owns survives with it: a reattached kernel's running job keeps
its heavy-job slot, and one found dead after a restart has its jobs ended.
"""

import json
import os
import signal
import subprocess
import time
import unittest

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


def bridges(home):
    """Pids of the bridges whose kernel lives under this daemon's home."""
    listing = subprocess.run(
        ["ps", "-axo", "pid=,command="], capture_output=True, text=True, check=True
    ).stdout
    return [
        int(line.split(None, 1)[0])
        for line in listing.splitlines()
        if "albedo_bridge.py" in line and str(home) in line
    ]


def alive(pid):
    try:
        os.kill(pid, 0)
    except ProcessLookupError:
        return False
    return True


@exclusive
class KernelDetachTests(unittest.TestCase):
    def setUp(self):
        self.cells = []

        def script(request):
            if request["input"][-1].get("role") == "user" and self.cells:
                return python(self.cells.pop(0))
            return text("ok")

        self.provider = Provider(script)
        self.addCleanup(self.provider.close)

        def prepare(app):
            # One heavy slot, granted to a job half a second in.
            app.daemon.env.update(
                ALBEDO_MAX_LOCAL_JOBS="1",
                ALBEDO_JOB_LOAD="0",
                ALBEDO_JOB_GRACE_SECONDS="0.5",
            )

        self.app = Albedo(self.provider, protocol="responses", prepare=prepare)
        self.app.__enter__()
        self.addCleanup(self.app.__exit__, None, None, None)
        self.session = self.app.session()

    def cell(self, code, *, during=None):
        """Run one cell through a model turn and return its tool result."""
        self.cells.append(code)
        self.app.prompt(self.session, "run it").close()
        if during:
            during()
        self.app.idle(self.session)
        results = [
            json.loads(event["result"])
            for event in self.app.events(self.session)
            if event.get("type") == "tool" and event.get("name") == "python"
        ]
        return results[-1]

    def test_variables_survive_graceful_and_crash_restarts(self):
        first = self.cell("import os\nsurvivor = 41\nos.getpid()")
        self.app.restart()
        after_stop = self.cell("survivor += 1\n(survivor, os.getpid())")
        self.assertEqual(after_stop["value"], f"(42, {first['value']})")
        self.app.restart(crash=True)
        after_crash = self.cell("survivor += 1\n(survivor, os.getpid())")
        self.assertEqual(after_crash["value"], f"(43, {first['value']})")
        notes = [
            event["text"]
            for event in self.app.events(self.session)
            if event.get("type") == "note"
        ]
        self.assertFalse([note for note in notes if "variables are gone" in note])

    def test_a_bridge_killed_mid_cell_still_delivers_the_result(self):
        first = self.cell("import os\nos.getpid()")
        (bridge,) = bridges(self.app.home)

        def kill_bridge():
            time.sleep(1)
            os.kill(bridge, signal.SIGKILL)

        slept = self.cell(
            "import time\ntime.sleep(3)\nprint('woke')\nos.getpid()",
            during=kill_bridge,
        )
        self.assertEqual(slept["status"], "ok", slept)
        self.assertEqual(slept["output"].strip(), "woke")
        self.assertEqual(slept["value"], first["value"])
        (replacement,) = bridges(self.app.home)
        self.assertNotEqual(replacement, bridge)

    def test_a_reattached_kernels_running_job_keeps_its_slot(self):
        held = self.cell(
            "import asyncio\na = run('sleep', '60')\nawait asyncio.sleep(1.5)\na.queued"
        )
        self.assertEqual(held["value"], "False")
        self.app.restart()
        # The new daemon's pool counts the slot a's job already holds.
        waiting = self.cell(
            "b = run('sleep', '60')\nawait asyncio.sleep(1.5)\n(a.queued, b.queued)"
        )
        self.assertEqual(waiting["value"], "(False, True)")
        freed = self.cell("await a.stop()\nawait asyncio.sleep(1)\nb.queued")
        self.assertEqual(freed["value"], "False")
        self.cell("await b.stop()")

    def test_a_kernel_killed_while_the_daemon_was_down_has_its_jobs_ended(self):
        started = self.cell(
            "import asyncio, os\nj = run('sleep', '300')\nawait asyncio.sleep(0.5)\n(os.getpid(), j.process.pid)"
        )
        kernel, job = (int(part) for part in started["value"].strip("()").split(","))

        # Graceful stop leaves the kernel running; it dies hard before the
        # next daemon reaches it, and its job's group outlives it.
        self.app.restart(prepare=lambda _: os.kill(kernel, signal.SIGKILL))
        deadline = time.monotonic() + 15
        while alive(job) and time.monotonic() < deadline:
            time.sleep(0.1)
        self.assertFalse(alive(job), "the dead kernel's job is still running")


if __name__ == "__main__":
    unittest.main()