"""Fresh spare kernels are single-use, bounded, and never adopted across a bundle edit."""

import ast
import json
import os
from pathlib import Path
import socket
import subprocess
import time
import unittest
import urllib.parse

from harness import Albedo, Provider, alive, exclusive, python, text, wait_until
from inspect_support import Inspection
from kernel_support import MARKER, ModelCells, editable_daemon

# An ssh that runs its command here with no environment of its own. Host
# `oldhost` reaches a python too old for a kernel, so remote.connect degrades
# to command mode: a connection that holds no process or descriptor.
FAKE_SSH = """#!/bin/sh
while [ $# -gt 0 ]; do
  case "$1" in
    -o) shift 2 ;;
    -*) shift ;;
    *) break ;;
  esac
done
target=$1
shift
path={path}
[ "$target" = oldhost ] && path={old}:$path
exec env -i HOME={home} PATH="$path" SHELL={old}/loginsh /bin/sh -c "$*"
"""


@exclusive
class KernelPoolTests(ModelCells, unittest.TestCase):
    bundle: Path

    def boot(self, *, spares=1, release=False, ssh=False, zygote=None):
        self.cells = []

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

        def prepare(app):
            launcher, self.bundle = editable_daemon(app.root)
            app.env.update(
                ALBEDO_DAEMON=str(launcher),
                ALBEDO_INSPECT="1",
                ALBEDO_KERNEL_SPARES=str(spares),
                ALBEDO_KERNEL_RELEASE_WITH_SPARE="1" if release else "0",
            )
            if zygote is not None:
                app.env["ALBEDO_KERNEL_ZYGOTE"] = zygote
            if ssh:
                # On the daemon's PATH, so a kernel's environment is unchanged
                # by connecting and only the connection itself can keep it.
                local, old, home = (app.root / name for name in ("bin", "old", "rhome"))
                for directory in (local, old, home):
                    directory.mkdir()
                path = f"{app.daemon.env['PATH']}"
                (local / "ssh").write_text(
                    FAKE_SSH.format(path=path, old=old, home=home)
                )
                (old / "python3").write_text("#!/bin/sh\necho 3.9\n")
                # A login shell that sources no profile, which could reset PATH.
                (old / "loginsh").write_text(
                    '#!/bin/sh\n[ "$1" = -l ] && shift\nexec /bin/sh "$@"\n'
                )
                for script in (local / "ssh", old / "python3", old / "loginsh"):
                    script.chmod(0o755)
                app.daemon.env["PATH"] = f"{local}{os.pathsep}{path}"

        self.provider = Provider(script)
        self.addCleanup(self.provider.close)
        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()
        self.inspection = Inspection(self.app)

    def spare_pids(self):
        """Pids of ready spares: the runtime's live spare custodians label
        themselves with their process. An adopted kernel's custodian is gone."""
        listed = self.inspection.evaluate(
            "io:write([Call(erlang, binary_to_integer, [Id]) || P <- Call(erlang, processes, []), "
            "{albedo_spare, Id} <- [Call(proc_lib, get_label, [P])]])"
        )
        return {int(pid) for pid in ast.literal_eval(listed or "[]") if alive(int(pid))}

    def ready(self, *, excluding=()):
        found = set()

        def available():
            nonlocal found
            found = self.spare_pids() - set(excluding)
            if not found:
                return False
            return (
                self.inspection.evaluate(
                    "Names = Call(erlang, registered, []), "
                    '[Name] = [N || N <- Names, lists:prefix("albedo_registry", atom_to_list(N))], '
                    "Subject = Call('gleam@erlang@process', named_subject, [Name]), "
                    "{ok, Runtime} = Call('albedo@daemon@registry', host, [Subject]), "
                    "io:write(Call('albedo@harness@runtime', spare_available, [Runtime]))"
                )
                == "true"
            )

        wait_until(available, 15, "the runtime did not refill its fresh spare")
        return found

    def value(self, code):
        result, _ = self.cell(code)
        self.assertEqual(result["status"], "ok", result)
        return ast.literal_eval(result["value"])

    def test_changed_environment_keeps_its_kernel_without_saving_secrets(self):
        self.boot(release=True)
        self.ready()
        pid = self.value(
            "import os\nos.environ['DENSITY_EPHEMERAL_SECRET']='private'\nos.getpid()"
        )
        time.sleep(1.2)
        self.assertTrue(alive(pid), "environment changes were lost at turn end")
        self.assertEqual(
            self.value("(os.getpid(), os.environ['DENSITY_EPHEMERAL_SECRET'])"),
            (pid, "private"),
        )
        state = self.app.home / "kernels" / f"{self.session}.state"
        self.assertFalse(state.exists(), "environment values must not be saved")
        self.value("del os.environ['DENSITY_EPHEMERAL_SECRET']\n42")
        wait_until(
            lambda: not alive(pid), 15, "restoring the environment did not unpin"
        )

    def test_user_thread_survives_turn_end(self):
        self.boot(release=True)
        self.ready()
        pid = self.value(
            "import os, threading\n_gate=threading.Event()\n_worker=threading.Thread(target=_gate.wait, daemon=True)\n_worker.start()\nos.getpid()"
        )
        time.sleep(1.2)
        self.assertTrue(alive(pid), "a running user thread was terminated")
        self.assertEqual(self.value("(os.getpid(), _worker.is_alive())"), (pid, True))
        self.value("_gate.set()\n_worker.join()\ndel _gate, _worker\n42")
        wait_until(
            lambda: not alive(pid), 15, "a finished user thread pinned its kernel"
        )

    def test_listening_socket_survives_turn_end(self):
        self.boot(release=True)
        self.ready()
        pid, port = self.value(
            "import os, socket\n_server=socket.socket()\n_server.bind(('127.0.0.1', 0))\n_server.listen()\n(os.getpid(), _server.getsockname()[1])"
        )
        time.sleep(1.2)
        self.assertTrue(alive(pid), "a listening server was terminated")
        with socket.create_connection(("127.0.0.1", port), timeout=2):
            pass
        self.assertEqual(self.value("os.getpid()"), pid)
        self.value("_server.close()\ndel _server\n42")
        wait_until(lambda: not alive(pid), 15, "a closed server pinned its kernel")

    def test_user_timer_survives_turn_end(self):
        self.boot(release=True)
        self.ready()
        pid = self.value(
            "import os\n_timer=asyncio.get_running_loop().call_later(300, lambda: None)\nos.getpid()"
        )
        time.sleep(1.2)
        self.assertTrue(alive(pid), "a scheduled user callback was discarded")
        self.assertEqual(self.value("os.getpid()"), pid)
        self.value("_timer.cancel()\ndel _timer\n42")
        wait_until(
            lambda: not alive(pid), 15, "a cancelled user timer pinned its kernel"
        )

    def test_user_async_task_survives_turn_end(self):
        self.boot(release=True)
        self.ready()
        pid = self.value(
            "import os\n_task=asyncio.create_task(asyncio.sleep(300))\nos.getpid()"
        )
        time.sleep(1.2)
        self.assertTrue(alive(pid), "a user asyncio task was cancelled")
        self.assertEqual(self.value("(os.getpid(), _task.done())"), (pid, False))
        self.value(
            "_task.cancel()\nawait asyncio.gather(_task, return_exceptions=True)\ndel _task\n42"
        )
        wait_until(lambda: not alive(pid), 15, "a finished user task pinned its kernel")

    def test_owned_cua_driver_survives_turn_end(self):
        from cua_test import FAKE_DRIVER

        self.boot(release=True)
        route = f"/sessions/{self.session}?view=configuration"
        with self.app.api(route) as response:
            response.read()
            etag = response.headers["ETag"]
        self.app.api(
            route,
            {"selection": {"extensions": {"cua": True}}},
            method="PATCH",
            headers={"If-Match": etag},
        ).close()
        self.app.api(f"/sessions/{self.session}/reload", {"target": "session"}).close()
        self.ready()
        driver = self.app.workspace / "fake-cua-driver"
        driver.write_text(FAKE_DRIVER)
        driver.chmod(0o755)
        pid = self.value(
            f"cua.configure(executable={str(driver)!r}, socket='/tmp/fake-density-cua.sock')\nawait cua.health()\nos.getpid()"
        )
        time.sleep(1.2)
        self.assertTrue(alive(pid), "an owned desktop driver was ended")
        self.assertEqual(self.value("await cua.health()\nos.getpid()"), pid)
        self.value("await cua.close()\n42")
        wait_until(
            lambda: not alive(pid), 15, "a closed desktop driver pinned its kernel"
        )

    def test_owned_browser_survives_turn_end(self):
        self.boot(release=True)
        self.ready()
        pid = self.value(
            "_browser=await browser.spawn()\n_page=await _browser.new_page()\nos.getpid()"
        )
        time.sleep(1.2)
        self.assertTrue(alive(pid), "an owned browser was ended")
        self.assertEqual(self.value("await _page.evaluate('1+1')"), 2)
        self.value("await _browser.close()\ndel _page, _browser\n42")
        wait_until(lambda: not alive(pid), 15, "a closed browser pinned its kernel")

    def test_open_remote_connection_survives_turn_end(self):
        self.boot(release=True, ssh=True)
        self.ready()
        pid = self.value(
            "_rem = await remote.connect('oldhost')\nassert _rem._stdin is None\nos.getpid()"
        )
        time.sleep(1.2)
        self.assertTrue(alive(pid), "an open remote connection was ended")
        self.assertEqual(
            self.value("(os.getpid(), len(remote.connections()))"), (pid, 1)
        )
        self.value("await _rem.close()\ndel _rem\n42")
        wait_until(
            lambda: not alive(pid), 15, "a closed remote connection pinned its kernel"
        )

    def jobs_api(self, path, **query):
        suffix = ("?" + urllib.parse.urlencode(query)) if query else ""
        route = f"/extensions/run/sessions/{self.session}/jobs/{path}{suffix}"
        with self.app.api(route) as response:
            return json.load(response)

    def test_release_keeps_retained_output_and_the_job_book(self):
        self.boot(release=True)
        self.ready()
        # A cell whose output spills past the preview (its result is stored
        # by reference, so it is not read back here), then one naming it.
        self.cells.append("print('é' * 70000 + 'end')")
        self.app.prompt(self.session, "run it").close()
        self.app.idle(self.session)
        cell, pid, job_id, text = self.value(
            "import os, sys\n"
            "job = run(sys.executable, '-c', 'print(\"first \\\\u00e9\"); print(\"second\")')\n"
            "await job\n"
            "([c['id'] for c in output.list() if c['bytes'] > 100000][0], os.getpid(), job.id, job.read())"
        )
        self.assertEqual(text, "first é\nsecond\n")
        wait_until(lambda: not alive(pid), 15, "the resting kernel was not released")
        # No kernel now: the routes answer from the saved book and wake none.
        [job] = self.jobs_api("recent")["items"]
        self.assertEqual(
            (job["id"], job["state"], job["exit_code"]), (job_id, "exited", 0)
        )
        self.assertEqual(job["bytes"], 16)
        page = self.jobs_api(f"{job_id}/output", offset=6, limit=4)
        self.assertEqual(
            (page["offset"], page["next_offset"], page["text"]), (6, 10, "é\ns")
        )
        self.assertEqual(self.jobs_api(f"{job_id}/output")["text"], text)
        self.assertEqual(self.spare_pids() & {pid}, set())
        restored = self.value(
            f"(output.read({cell!r}, offset=139990, limit=100), jobs[{job_id!r}].tail(lines=1), "
            f"jobs[{job_id!r}].exit_code, job.read())"
        )
        self.assertEqual(restored, ("ééééé" + "end\n", "second\n", 0, text))

    def test_spare_is_adopted_released_once_and_refilled(self):
        self.boot(release=True)
        before = self.ready()
        first = self.value("import os\nsentinel = 41\nos.getpid()")
        self.assertIn(first, before, "the first cell did not adopt a spare")
        wait_until(lambda: not alive(first), 15, "the resting kernel was not released")
        next_spares = self.ready()
        self.assertTrue(before.isdisjoint(next_spares))
        second, sentinel = self.value("(os.getpid(), sentinel)")
        self.assertIn(second, next_spares)
        self.assertNotEqual(first, second)
        self.assertEqual(sentinel, 41)
        wait_until(
            lambda: not alive(second), 15, "the adopted kernel returned to the pool"
        )
        self.assertTrue(self.ready().isdisjoint({first, second}))

    def test_live_background_work_keeps_the_adopted_kernel(self):
        self.boot(release=True)
        self.ready()
        gate = self.app.workspace / "background-gate"
        program = f"import pathlib, time; p = pathlib.Path({str(gate)!r}); exec('while not p.exists(): time.sleep(0.01)')"
        pid = self.value(
            f"import os, sys\njob = run(sys.executable, '-c', {program!r})\nos.getpid()"
        )
        self.assertTrue(alive(pid))
        self.assertEqual(self.value("(os.getpid(), job.exit_code)"), (pid, None))
        gate.touch()
        self.value("await job\n42")
        wait_until(lambda: not alive(pid), 15, "a finished job pinned its kernel")

    def zygote_of(self, pid):
        """The zygote that forked this spare, or None."""
        parent = subprocess.check_output(
            ["ps", "-o", "ppid=", "-p", str(pid)], text=True
        ).strip()
        command = subprocess.run(
            ["ps", "-o", "command=", "-p", parent], capture_output=True, text=True
        ).stdout
        return int(parent) if "albedo_zygote.py" in command else None

    def test_zygote_forks_spares_and_is_recycled_after_a_bundle_edit(self):
        self.boot(zygote="1")
        [spare] = self.ready()
        zygote = self.zygote_of(spare)
        self.assertIsNotNone(zygote, "the spare was not forked by a zygote")
        self.assertEqual(self.value("import os\nos.getpid()"), spare)
        with (self.bundle / "albedo_bundle.py").open("a") as source:
            source.write(MARKER)
        # A new session's kernel must run the edited code, not the zygote's.
        self.session = self.app.session()
        self.assertEqual(
            self.value("import albedo_bundle\nalbedo_bundle.MARKER"), "new bundle"
        )
        wait_until(
            lambda: not alive(zygote), 15, "the stale zygote outlived its bundle"
        )

    def test_opens_past_the_spares_fork_from_the_zygote(self):
        self.boot(zygote="1")
        [spare] = self.ready()
        zygote = self.zygote_of(spare)
        sessions = [self.app.session() for _ in range(3)]
        for session in sessions:
            self.cells.append("import os\nos.getpid()")
            self.app.prompt(session, "run it").close()
        pids = []
        for session in sessions:
            self.session = session
            self.app.idle(session)
            [result] = [
                json.loads(part["value"])
                for entry in self.app.history(session)["items"]
                if entry["kind"] == "tool_result"
                for part in entry["content"]
                if part["kind"] == "json" and part["field"] == "result"
            ]
            pids.append(ast.literal_eval(result["value"]))
        self.assertEqual({self.zygote_of(pid) for pid in pids}, {zygote})

    def test_an_empty_pool_uses_a_cold_kernel(self):
        self.boot(spares=0)
        pid = self.value("import os\nos.getpid()")
        self.assertTrue(alive(pid))
        self.assertEqual(self.spare_pids(), set())

    def test_a_changed_bundle_is_not_adopted(self):
        self.boot()
        before = self.ready()
        with (self.bundle / "albedo_bundle.py").open("a") as source:
            source.write(MARKER)
        pid, marker = self.value(
            "import os, albedo_bundle\n(os.getpid(), albedo_bundle.MARKER)"
        )
        self.assertNotIn(pid, before)
        self.assertEqual(marker, "new bundle")
        wait_until(
            lambda: all(not alive(old) for old in before),
            15,
            "the stale spare survived its rejected adoption",
        )

    def test_shutdown_stops_spares_but_detaches_an_adopted_kernel(self):
        self.boot()
        self.ready()
        adopted = self.value("import os\nsentinel = 41\nos.getpid()")
        ready = self.ready(excluding={adopted})
        self.app.restart()
        wait_until(
            lambda: all(not alive(pid) for pid in ready),
            15,
            "a fresh spare survived daemon shutdown",
        )
        self.assertTrue(alive(adopted), "shutdown ended an adopted detached kernel")
        self.assertEqual(
            self.value("(os.getpid(), sentinel)"),
            (adopted, 41),
            "the restarted daemon did not reattach the adopted kernel",
        )