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

The kernel runs detached and the daemon connects to its socket directly; it
reattaches after a restart (graceful or crash) and after the connection drops
mid-cell, so the namespace, the kernel process, and an in-flight cell's result
all survive. A kernel found dead after a restart has its jobs ended.
"""

import ast
import json
import os
import shutil
import signal
import sqlite3
import subprocess
from contextlib import closing, suppress
from pathlib import Path
import time
import unittest

import harness
from harness import alive
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 kill(*pids):
    for pid in pids:
        with suppress(OSError):
            os.kill(pid, signal.SIGKILL)


# A live job keeps an idle kernel from being released between turns, as a
# user's running work does.
HOLD = "hold = run('sleep', '120')\n"


# exclusive: restarts the daemon
@exclusive
class KernelDetachTests(unittest.TestCase):
    def setUp(self):
        self.cells = []
        self.cell_timeout_ms = 60_000

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

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

        self.app = Albedo(
            self.provider,
            protocol="responses",
            prepare=lambda app: app.env.update(
                ALBEDO_TEST_TIME_SCALE="5", ALBEDO_IDLE_SECONDS="600"
            ),
        )
        self.app.__enter__()
        self.addCleanup(self.app.__exit__, None, None, None)
        self.session = self.app.session()

    def cell(self, code):
        """Run one cell through a model turn and return its tool result."""
        self.cells.append(code)
        self.app.prompt(self.session, "run it").close()
        self.app.idle(self.session)
        results = [
            json.loads(part["value"])
            for entry in self.app.history(self.session)["items"]
            if entry["kind"] == "tool_result" and entry["tool"]["name"] == "python"
            for part in entry["content"]
            if part["kind"] == "json" and part["field"] == "result"
        ]
        return results[-1]

    @exclusive
    def test_variables_survive_graceful_and_crash_restarts(self):
        first = self.cell(HOLD + "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 = [
            part["text"]
            for entry in self.app.history(self.session)["items"]
            if entry["kind"] == "note"
            for part in entry["content"]
            if part["kind"] == "text"
        ]
        self.assertFalse([note for note in notes if "variables are gone" in note])

    @exclusive
    def test_a_kernel_working_in_an_albedo_checkout_reattaches_after_a_restart(self):
        """A workspace with its own, different priv/python (another albedo
        checkout) is not the bundle the kernel runs or is checked against."""
        checkout = Path(self.app.workspace) / "checkout"
        shutil.copytree(
            Path(__file__).resolve().parents[2] / "priv" / "python",
            checkout / "priv" / "python",
            ignore=shutil.ignore_patterns("__pycache__"),
        )
        with (checkout / "priv" / "python" / "albedo_bundle.py").open("a") as file:
            file.write("# another checkout's edit\n")
        self.session = self.app.session(checkout)
        first = self.cell(HOLD + "import os\nsurvivor = 41\nos.getpid()")
        self.app.restart()
        after = self.cell("survivor += 1\n(survivor, os.getpid())")
        self.assertEqual(after["value"], f"(42, {first['value']})")

    def test_stopping_with_kernels_ends_them_instead_of_leaving_them(self):
        kernel = int(self.cell("import os\nsurvivor = 41\nos.getpid()")["value"])
        self.app.cli("daemon", "--stop", "--kernels")
        harness.wait_until(
            lambda: not alive(kernel), 15, "the kernel outlived daemon --stop --kernels"
        )
        self.app.restart()
        # A fresh kernel, with the variables the closing session saved.
        after = self.cell("survivor += 1\n(survivor, os.getpid())")
        survivor, pid = ast.literal_eval(after["value"])
        self.assertEqual(survivor, 42)
        self.assertNotEqual(pid, kernel)

    @exclusive
    def test_a_daemon_booted_on_a_copy_of_the_database_leaves_its_kernels_alone(self):
        """The copy carries this daemon's kernel records (socket, token,
        process group); its daemon forgets them instead of attaching."""
        first = self.cell(HOLD + "import os\nsurvivor = 41\nos.getpid()")
        copy = harness.Daemon("copy")
        with (
            closing(sqlite3.connect(self.app.home / "albedo.sqlite")) as source,
            closing(sqlite3.connect(copy.home / "albedo.sqlite")) as target,
        ):
            source.backup(target)
        copy.boot()
        self.addCleanup(copy._stop, 15)
        with closing(sqlite3.connect(copy.home / "albedo.sqlite")) as db:
            (left,) = db.execute("SELECT count(*) FROM kernel_links").fetchone()
        self.assertEqual(left, 0)
        self.assertIn(
            "kernel records: left 1 kernels of another daemon's home alone",
            (copy.home / "daemon.log").read_text(),
        )
        after = self.cell("survivor += 1\n(survivor, os.getpid())")
        self.assertEqual(after["value"], f"(42, {first['value']})")

    @exclusive
    def test_a_kernel_unreachable_at_restart_is_kept_until_it_answers(self):
        """Only proof of absence forgets a kernel: a socket that refuses the
        attach for another reason leaves the kernel, its namespace and its
        jobs alone, and the daemon attaches once the socket answers."""
        first = self.cell(HOLD + "import os\nsurvivor = 41\nos.getpid()")
        kernel = int(first["value"])
        with closing(sqlite3.connect(self.app.home / "albedo.sqlite")) as db:
            (run_dir,) = db.execute(
                "SELECT run_dir FROM kernel_links WHERE session=?", (self.session,)
            ).fetchone()
        socket_path = Path(run_dir) / "kernel.sock"
        # connect(2) fails with EACCES, which says nothing about the kernel.
        self.app.restart(prepare=lambda app: socket_path.chmod(0))
        log = self.app.home / "daemon.log"
        harness.wait_until(
            lambda: (
                not alive(kernel) or "kernel_unreachable_at_attach" in log.read_text()
            ),
            15,
            "the daemon never tried to attach",
        )
        self.assertTrue(alive(kernel), "an unreachable kernel was ended")
        socket_path.chmod(0o600)
        after = self.cell("survivor += 1\n(survivor, os.getpid())")
        self.assertEqual(after["value"], f"(42, {kernel})")

    def orphaned_and_reaped(self, unreachable):
        """Restart with the kernel made unreachable by `unreachable(pid,
        run_dir)` and check the whole fallback: the old kernel is orphaned
        once its tries run out, the session gets a fresh kernel that tells the
        model what it lost, and the old kernel and its job are ended at the
        idle limit."""
        started = self.cell(
            HOLD + "import asyncio, os\nawait asyncio.sleep(0.3)\n"
            "(os.getpid(), hold.process.pid, hold.id)"
        )
        kernel, job, handle = ast.literal_eval(started["value"])
        # A stopped kernel never ends by itself: never leave one behind.
        self.addCleanup(kill, kernel, job)
        with closing(sqlite3.connect(self.app.home / "albedo.sqlite")) as db:
            kernel_id, run_dir = db.execute(
                "SELECT kernel, run_dir FROM kernel_links WHERE session=?",
                (self.session,),
            ).fetchone()

        def prepare(app):
            unreachable(kernel, Path(run_dir))
            # 40 tries take about 2 s; the idle limit is 10 s after them.
            app.env.update(ALBEDO_TEST_TIME_SCALE="40", ALBEDO_IDLE_SECONDS="400")

        self.app.restart(prepare=prepare)
        log = self.app.home / "daemon.log"
        harness.wait_until(
            lambda: "{kernel_unreachable," in log.read_text(),
            30,
            "the daemon never gave up on the unreachable kernel",
        )

        self.cells.append("import os\nos.getpid()")
        self.app.prompt(self.session, "run it").close()
        self.app.idle(self.session)
        prompt = self.last_prompt()
        self.assertIn("The previous python kernel could not be reached", prompt)
        self.assertIn(
            f"Jobs left in the unreachable kernel: {handle} (sleep 120)", prompt
        )
        fresh = int(self.cell("import os\nos.getpid()")["value"])
        self.assertNotEqual(fresh, kernel)
        self.assertTrue(alive(kernel), "the orphan was ended before the idle limit")
        with closing(sqlite3.connect(self.app.home / "albedo.sqlite")) as db:
            orphans = db.execute("SELECT pid FROM kernel_orphans").fetchall()
        self.assertEqual(orphans, [(kernel,)])

        harness.wait_until(
            lambda: not alive(kernel) and not alive(job),
            30,
            "the orphaned kernel or its job outlived the idle limit",
        )
        harness.wait_until(
            lambda: f"kernel orphan {kernel_id}: ended" in log.read_text(),
            10,
            "the orphan was not reaped by the daemon",
        )
        self.assertFalse(Path(run_dir).exists())
        # The reap ended the orphan and nothing else: the session still runs
        # cells (its fresh kernel may have been released at rest since, which
        # is the pool's ordinary release, not the reap).
        reaps = [
            line for line in log.read_text().splitlines() if "kernel orphan " in line
        ]
        self.assertEqual(reaps, [f"kernel orphan {kernel_id}: ended"])
        after = self.cell("import os\nos.getpid()")
        self.assertEqual(after["status"], "ok", after)
        self.assertNotEqual(int(after["value"]), kernel)

    @exclusive
    def test_a_kernel_unreachable_for_every_try_is_replaced_and_reaped(self):
        # connect(2) fails with EACCES on every try.
        self.orphaned_and_reaped(lambda _, run_dir: (run_dir / "kernel.sock").chmod(0))

    @exclusive
    def test_a_stopped_kernel_is_replaced_and_reaped(self):
        # The socket accepts, but a stopped kernel never says hello.
        self.orphaned_and_reaped(lambda pid, _: os.kill(pid, signal.SIGSTOP))

    @exclusive
    def test_a_busy_kernel_at_restart_is_waited_for_not_orphaned(self):
        """A kernel holding the GIL in C cannot say hello, but it is neither
        stopped nor gone: the daemon waits for it without spending tries,
        tells the person once that it is busy, and attaches to the same
        process and namespace when it answers."""
        gate = self.app.workspace / "busy-gate"
        first = self.cell(
            "import ctypes, os, threading, time\n"
            "def hold():\n"
            f"    while not os.path.exists({str(gate)!r}):\n"
            "        time.sleep(0.05)\n"
            "    ctypes.PyDLL(None).sleep(8)  # libc sleep, GIL held\n"
            "threading.Thread(target=hold, daemon=True).start()\n"
            "survivor = 41\nos.getpid()"
        )
        kernel = int(first["value"])

        def prepare(app):
            gate.touch()
            # The 40 tries would take about 2 s; the GIL is held for 8.
            app.env.update(ALBEDO_TEST_TIME_SCALE="40", ALBEDO_IDLE_SECONDS="400")

        self.app.restart(prepare=prepare)
        cursor = self.app.stream_page(self.session)
        # Scaled to 90 s: the cell waits for the kernel to answer.
        self.cell_timeout_ms = 3_600_000
        after = self.cell("survivor += 1\n(survivor, os.getpid())")
        self.assertEqual(after["value"], f"(42, {kernel})")
        notes = [
            event["data"]["text"]
            for event in self.app.stream_page(self.session, cursor)["events"]
            if event.get("type") == "note"
        ]
        busy = (
            "the python kernel is busy and has not answered since the daemon restarted"
        )
        self.assertEqual(len([note for note in notes if busy in note]), 1, notes)
        log = (self.app.home / "daemon.log").read_text()
        self.assertIn("kernel_busy_at_attach", log)
        self.assertNotIn("{kernel_unreachable,", log)
        with closing(sqlite3.connect(self.app.home / "albedo.sqlite")) as db:
            self.assertEqual(
                db.execute("SELECT count(*) FROM kernel_orphans").fetchone(), (0,)
            )

    @exclusive
    def test_a_connection_dropped_mid_cell_still_delivers_the_result(self):
        first = self.cell(HOLD + "import os\nos.getpid()")
        self.assertEqual(bridges(self.app.home), [], "a local kernel kept a bridge")
        # The kernel drops the daemon's connection from its side a second in.
        slept = self.cell(
            "import socket, sys, threading\n"
            "link = sys.modules['__main__'].LINK\n"
            "threading.Timer(1, link.connection.shutdown, [socket.SHUT_RDWR]).start()\n"
            "threading.Event().wait(3)\nprint('woke')\nos.getpid()"
        )
        self.assertEqual(slept["status"], "ok", slept)
        self.assertEqual(slept["output"].strip(), "woke")
        self.assertEqual(slept["value"], first["value"])
        self.assertEqual(bridges(self.app.home), [])

    def test_concurrent_jobs_keep_running_and_deadlines_count_wall_time(self):
        # Legacy admission settings must not throttle new kernels.
        self.app.daemon.env.update(
            ALBEDO_MAX_LOCAL_JOBS="1", ALBEDO_JOB_GRACE_SECONDS="0.2"
        )
        self.app.restart()
        program = "import time; time.sleep(5.5); print('progress', flush=True); time.sleep(60)"
        # The three jobs each use a seven-second wall timeout; give the cell
        # sixty real seconds after its fivefold timer scaling for spawn and drain.
        self.cell_timeout_ms = 300_000
        result = self.cell(
            "import sys\n"
            f"workers = [run(sys.executable, '-c', {program!r}, timeout=7) for _ in range(3)]\n"
            "for worker in workers:\n    await worker\n"
            "[(worker.timed_out, worker.duration, worker.tail().startswith('progress')) for worker in workers]"
        )
        self.assertEqual(result["status"], "ok", result)
        outcomes = ast.literal_eval(result["value"])
        self.assertEqual(len(outcomes), 3)
        for timed_out, duration, progressed in outcomes:
            self.assertTrue(timed_out)
            self.assertTrue(progressed)
            self.assertGreaterEqual(duration, 7)
            self.assertLess(duration, 10)

    def last_prompt(self):
        """The newest model request's user text: where a notice would ride."""
        request = next(
            record["request"]
            for record in reversed(self.provider.requests)
            if any(
                item.get("role") == "user"
                for item in record["request"].get("input", [])
            )
        )
        return [
            item["content"] for item in request["input"] if item.get("role") == "user"
        ][-1]

    def test_a_restart_tells_the_model_what_the_kernel_is_still_running(self):
        started = self.cell(
            "import asyncio, os\nj = run('sleep', '120')\n"
            "await asyncio.sleep(0.3)\n(os.getpid(), j.id)"
        )
        kernel, job = (
            part.strip(" '\"") for part in started["value"].strip("()").split(",")
        )
        self.app.restart()

        # The next turn runs on the reattached kernel, and its request says
        # the kernel carried on and names the job still running, so a live
        # handle is not mistaken for one the restart killed.
        self.cells.append("import os\nos.getpid()")
        self.app.prompt(self.session, "run it").close()
        self.app.idle(self.session)
        prompt = self.last_prompt()
        self.assertIn("<system-note>The python kernel reattached", prompt)
        self.assertIn("Jobs still running: " + job, prompt)
        self.assertIn("sleep 120", prompt)
        self.assertEqual(self.cell("import os\nos.getpid()")["value"], str(int(kernel)))
        # The notice rides the first request after the reattach, once.
        self.cells.append("import os\nos.getpid()")
        self.app.prompt(self.session, "run it").close()
        self.app.idle(self.session)
        self.assertNotIn("The python kernel reattached", self.last_prompt())

    def test_a_job_ending_after_a_restart_wakes_its_idle_session(self):
        started = self.cell("j = run('sleep', '4')\nj.id")
        job = started["value"].strip("'\"")
        requests = len(self.provider.requests)
        # The restart leaves the session idle and unloaded; only its kernel
        # comes back, and the job's end must still reach the model.
        self.app.restart()
        deadline = time.monotonic() + 30
        while len(self.provider.requests) == requests and time.monotonic() < deadline:
            time.sleep(0.1)
        self.assertGreater(len(self.provider.requests), requests, "no wake")
        self.app.idle(self.session)
        self.assertIn(job, self.last_prompt())

    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, j.id)"
        )
        kernel, job, handle = ast.literal_eval(started["value"])

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

        # The new kernel's first request tells the model what died with the
        # old one, from the job record the daemon kept.
        self.cells.append("import os\nos.getpid()")
        self.app.prompt(self.session, "run it").close()
        self.app.idle(self.session)
        prompt = self.last_prompt()
        self.assertIn("<system-note>The python kernel restarted", prompt)
        self.assertIn(f"Jobs stopped by the restart: {handle} (sleep 300)", prompt)

    @exclusive
    def test_a_restart_tells_the_model_when_nothing_was_running(self):
        self.cell("survivor = 1")
        self.app.restart()
        self.cells.append("survivor")
        self.app.prompt(self.session, "run it").close()
        self.app.idle(self.session)
        self.assertIn(
            "The python kernel reattached; its variables and jobs carried on."
            " Nothing was running.",
            self.last_prompt(),
        )


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