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

import ast
import os
from pathlib import Path
import subprocess
import socket
import time
import unittest

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):
        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 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: children of a waiting `albedo_bridge.py spare`.

        An adopted kernel keeps the spare script's name but no helper parent."""
        rows = [
            line.split(None, 2)
            for line in subprocess.check_output(
                ["ps", "-axo", "pid=,ppid=,command="], text=True
            ).splitlines()
        ]
        script = str(self.bundle / "albedo_spare.py")
        helper = str(self.bundle / "albedo_bridge.py") + " spare"
        helpers = {int(pid) for pid, _, command in rows if helper in command}
        return {
            int(pid)
            for pid, ppid, command in rows
            if script in command and int(ppid) in helpers
        }

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