"""Lazy kernels: a turn holds no Python kernel until its first tool call."""

import json
from pathlib import Path
import subprocess
import threading
import time
import unittest

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

IDLE_SECONDS = 2


def kernels(home):
    """Pids of the kernels whose run directory is 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_kernel.py" in line and f"{home}/run/" in line
    ]


def wait_until(predicate, timeout=20):
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if predicate():
            return True
        time.sleep(0.05)
    return predicate()


def inputs(record):
    return record["request"]["input"]


def is_note(item):
    return item.get("role") == "user" and str(item.get("content", "")).startswith(
        "<system-note>"
    )


def prompt_of(request):
    """The latest user prompt, not a system note that follows a tool result."""
    users = [
        item
        for item in request["input"]
        if item.get("role") == "user" and not is_note(item)
    ]
    return str(users[-1].get("content", "")).split("<system-note>")[0] if users else ""


# exclusive: lazy kernels and the idle timeout are daemon-wide, and the tests
# count every kernel the daemon runs
@exclusive
class KernelLazyTests(unittest.TestCase):
    def setUp(self):
        self.model_gate = threading.Event()
        self.model_waiting = threading.Event()
        self.cells = {}
        self.follow = {}

        def script(request):
            last = request["input"][-1]
            prompt = prompt_of(request)
            if last.get("role") == "user" and not is_note(last):
                if prompt in self.cells:
                    return python(self.cells[prompt])
                if prompt.startswith("wait for the model"):
                    self.model_waiting.set()
                    self.model_gate.wait(60)
                    return text("waited")
            outputs = [
                item
                for item in request["input"]
                if item.get("type") == "function_call_output"
            ]
            follow = self.follow.get(prompt, [])
            if prompt in self.cells and len(outputs) <= len(follow):
                # Each follow-up cell waits on the model first when asked to.
                wait, code = follow[len(outputs) - 1] if outputs else (False, None)
                if code is not None:
                    if wait:
                        self.model_waiting.set()
                        self.model_gate.wait(60)
                    return python(code)
            return text("ok")

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

    def start(self, slow_marker=None, crash_marker=None, **env):
        def prepare(app):
            # The spare pool stays off unless a test asks for it: the edited
            # kernel scripts below run only in a kernel booted cold, and the
            # idle reap counts kernel processes.
            app.env.update(
                {
                    "ALBEDO_TEST_TIME_SCALE": "100",
                    "ALBEDO_IDLE_SECONDS": str(IDLE_SECONDS),
                    "ALBEDO_KERNEL_SPARES": "0",
                    **env,
                }
            )
            if slow_marker is not None or crash_marker is not None:
                launcher, bundle = editable_daemon(app.root)
                kernel = bundle / "albedo_kernel.py"
                prelude = "import os as _os, sys as _sys, time as _time\n"
                if slow_marker is not None:
                    prelude += f"if _os.path.exists({str(slow_marker)!r}):\n    _time.sleep(3)\n"
                if crash_marker is not None:
                    prelude += f"if _os.path.exists({str(crash_marker)!r}):\n    _sys.exit(3)\n"
                kernel.write_text(prelude + kernel.read_text())
                app.env.update(ALBEDO_DAEMON=str(launcher))

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

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

    def test_text_turns_and_model_waits_hold_no_kernel(self):
        app = self.start()
        existing = set(kernels(app.home))
        session = app.session()
        self.turn(app, session, "just talk")
        self.assertEqual(set(kernels(app.home)) - existing, set())

        self.cells["remember"] = "answer = 7\nrows = [1, 2, 3]"
        self.turn(app, session, "remember")
        # The scaled idle window may already have released the kernel the cell
        # opened, so it is counted by the state its release saved.
        # Nothing ran before, so the first kernel has no namespace to call lost.
        self.assertFalse(
            any(is_note(item) for item in inputs(self.provider.requests[-1])),
            inputs(self.provider.requests[-1]),
        )
        self.assertTrue(
            wait_until(lambda: not set(kernels(app.home)) - existing),
            "idle kernel outlived its limit",
        )
        # The cell's kernel held the namespace: its release saved it.
        self.assertTrue((app.home / "kernels" / f"{session}.state").exists())

        app.prompt(session, "wait for the model").close()
        self.assertTrue(self.model_waiting.wait(30))
        time.sleep(0.5)
        self.assertEqual(
            set(kernels(app.home)) - existing,
            set(),
            "a turn waiting on its model must not hold a kernel",
        )
        self.model_gate.set()
        app.idle(session)

        self.cells["recall"] = "print('recalled', answer, rows)"
        self.turn(app, session, "recall")
        after = inputs(self.provider.requests[-1])
        self.assertTrue(
            any("recalled 7 [1, 2, 3]" in item.get("output", "") for item in after),
            after,
        )
        # The restore lost nothing, so the model is told nothing.
        notices = [
            json.dumps(item)
            for item in after
            if item.get("role") == "user" and "<system-note>" in json.dumps(item)
        ]
        self.assertEqual(notices, [], after)

    def test_a_restore_with_no_variables_says_nothing_unless_it_lost_one(self):
        app = self.start()
        existing = set(kernels(app.home))
        session = app.session()

        def released_then(prompt, code):
            self.assertTrue(
                wait_until(lambda: not set(kernels(app.home)) - existing),
                "idle kernel outlived its limit",
            )
            self.cells[prompt] = code
            self.turn(app, session, prompt)
            return [
                str(item.get("content"))
                for item in inputs(self.provider.requests[-1])
                if is_note(item)
            ]

        self.cells["print only"] = "print('nothing to keep')"
        self.turn(app, session, "print only")
        # Saved with no variables and nothing missed: a lossless restore.
        self.assertEqual(released_then("again", "print('again')"), [])
        # Saved with no variables but one it could not keep: still a loss.
        self.cells["unpicklable"] = "import threading\nlock = threading.Lock()"
        self.turn(app, session, "unpicklable")
        notes = released_then("after the lock", "print('after')")
        self.assertEqual(len(notes), 1, notes)
        self.assertIn("lock", notes[0])

    def test_interrupt_while_the_kernel_starts_runs_no_cell(self):
        slow = Path("/tmp") / f"albedo-slow-{time.monotonic_ns()}"
        self.addCleanup(lambda: slow.unlink(missing_ok=True))
        app = self.start(slow_marker=slow)
        session = app.session()
        written = app.workspace / "written"
        self.cells["write"] = f"open({str(written)!r}, 'w').write('x')"
        slow.touch()
        start = len(self.provider.requests)
        with app.prompt(session, "write") as response:
            input_id = json.load(response)["id"]
        self.assertTrue(wait_until(lambda: len(self.provider.requests) > start))
        with app.api(f"/sessions/{session}?tail=0") as response:
            captured = json.load(response)
        # The cell is now waiting on the slow kernel.
        time.sleep(0.5)
        interrupted = time.monotonic()
        app.api(
            f"/sessions/{session}/interrupt",
            {
                "run_id": captured["status"]["run_id"],
                "through_input_order": captured["input_order"],
            },
        ).close()

        def receipt():
            with app.api(f"/sessions/{session}/inputs/{input_id}") as response:
                return json.load(response).get("turn") or {}

        # The run ends at once, not when the kernel it was waiting on starts.
        self.assertTrue(
            wait_until(lambda: receipt().get("state") == "interrupted", timeout=2)
        )
        self.assertLess(time.monotonic() - interrupted, 2.0)
        app.idle(session)
        time.sleep(4)
        self.assertFalse(written.exists(), "a cancelled turn ran its cell")
        slow.unlink()
        self.cells["after"] = "print('after', 6 * 7)"
        self.turn(app, session, "after")
        self.assertTrue(
            any(
                "after 42" in item.get("output", "")
                for item in inputs(self.provider.requests[-1])
            )
        )

    def test_a_long_model_wait_mid_turn_releases_and_restores_the_kernel(self):
        # A zero memory budget is constant pressure, under which a long model
        # wait releases the kernel; the time scale makes the 3 s rest 30 ms.
        app = self.start(ALBEDO_KERNEL_BUDGET_MB="0", ALBEDO_KERNEL_SPARES="1")
        session = app.session()
        place = app.workspace / "place"
        place.mkdir()
        self.cells["work"] = f"import os\nanswer = 7\nos.chdir({str(place)!r})"
        self.follow["work"] = [
            (True, "import os\nprint('kept', answer, os.getcwd())"),
            (False, "print('again', answer)"),
        ]

        def kernel_state():
            with app.api(f"/sessions/{session}?tail=0") as response:
                return json.load(response)["kernel"]["state"]

        app.prompt(session, "work").close()
        self.assertTrue(self.model_waiting.wait(30))
        self.assertTrue(
            wait_until(lambda: kernel_state() == "none"),
            "a turn waiting on its model past the rest time kept its kernel",
        )
        self.model_gate.set()
        app.idle(session)
        outputs = [
            item.get("output", "")
            for item in inputs(self.provider.requests[-1])
            if item.get("type") == "function_call_output"
        ]
        self.assertTrue(any("kept 7" in output for output in outputs), outputs)
        self.assertTrue(
            any(
                f"{place.resolve()}" in output or str(place) in output
                for output in outputs
            ),
            outputs,
        )
        self.assertTrue(any("again 7" in output for output in outputs), outputs)

    def test_a_kernel_that_dies_while_starting_fails_only_its_call(self):
        crash = Path("/tmp") / f"albedo-crash-{time.monotonic_ns()}"
        self.addCleanup(lambda: crash.unlink(missing_ok=True))
        app = self.start(crash_marker=crash)
        session = app.session()
        self.cells["boom"] = "print('ran', 6 * 7)"
        crash.touch()
        self.turn(app, session, "boom")
        outputs = [
            item.get("output", "")
            for item in inputs(self.provider.requests[-1])
            if item.get("type") == "function_call_output"
        ]
        self.assertEqual(len(outputs), 1, outputs)
        self.assertIn("could not start the session python kernel", outputs[0])
        self.assertNotIn("ran 42", outputs[0])
        crash.unlink()
        self.cells["again"] = "print('ran', 6 * 7)"
        self.turn(app, session, "again")
        self.assertTrue(
            any(
                "ran 42" in item.get("output", "")
                for item in inputs(self.provider.requests[-1])
            )
        )

    def test_a_kernel_lost_mid_turn_is_replaced_at_the_next_call(self):
        app = self.start()
        session = app.session()
        self.cells["die"] = "import os; os._exit(1)"
        self.follow["die"] = [(False, "print('again', 6 * 7)")]
        self.turn(app, session, "die")
        request = inputs(self.provider.requests[-1])
        outputs = [
            item.get("output", "")
            for item in request
            if item.get("type") == "function_call_output"
        ]
        self.assertEqual(len(outputs), 2, outputs)
        self.assertIn("kernel lost", outputs[0])
        self.assertIn("again 42", outputs[1])
        self.assertTrue(
            any(
                "python kernel got reset" in str(item.get("content"))
                for item in request
            ),
            request,
        )


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