"""A restarted kernel gets back its helpers, not only its variables.

Saved state is written when an idle kernel is released and read when the next
turn opens a new one, so these scenarios run through the real daemon: the
restore path lives in the kernel and the notice in the daemon.
"""

import os
import pickle
import unittest

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

IDLE_SECONDS = 1
BIG_DEFINITION = "def big():\n    return '" + "x" * 40000 + "'\n"

DEFINE = (
    "import json\n"
    "import os.path\n"
    "import xml.dom\n"
    "def helper(x):\n    return x * 2\n"
    "class Box:\n    def __init__(self, v):\n        self.v = v\n"
    "answer = 7\n"
    "def rebound():\n    return 1\n"
    "if answer:\n    rebound = 5\n" + BIG_DEFINITION + "blob = bytes(9 * 1024 * 1024)\n"
)
FAIL = (
    "def survivor():\n    return 'alive'\n"
    "pass\n"
    "raise RuntimeError('boom')\n"
    "def never_ran():\n    return 'unreachable'\n"
)
RECALL = (
    "print('RESULT', helper(21), Box(3).v, json.dumps([1]), answer, rebound,\n"
    "      survivor(), os.path.basename('a/b'), xml.dom.__name__,\n"
    "      'big' in globals(), 'blob' in globals(), 'never_ran' in globals())"
)
LEGACY = "print('LEGACY', legacy)"


# exclusive: requires the daemon-wide idle-kernel timing knob
@exclusive
class KernelStateTests(unittest.TestCase):
    def setUp(self):
        cells = {"define": DEFINE, "fail": FAIL, "recall": RECALL, "legacy": LEGACY}

        def script(request):
            if request["input"][-1].get("role") == "user":
                prompt = latest_user(request).split("<system-note>")[0]
                if prompt in cells:
                    return python(cells[prompt])
            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_IDLE_SECONDS=str(IDLE_SECONDS)),
        )
        self.app.__enter__()
        self.addCleanup(self.app.__exit__, None, None, None)

    def state(self, session):
        return self.app.home / "kernels" / f"{session}.state"

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

    def saved_after(self, session, since):
        """The idle kernel was released again: its state file is newer than `since`."""
        wait_until(
            lambda: (
                self.state(session).exists()
                and self.state(session).stat().st_mtime_ns > since
            ),
            15,
            "the idle kernel was not saved",
        )

    def outputs(self):
        return "".join(
            item.get("output", "")
            for item in self.provider.requests[-1]["request"]["input"]
        )

    def test_helpers_classes_imports_and_failed_cells_survive_a_restart(self):
        session = self.app.session()
        self.turn(session, "define")
        self.saved_after(session, 0)
        saved = self.state(session).stat().st_mtime_ns
        self.turn(session, "fail")
        self.saved_after(session, saved)
        self.turn(session, "recall")
        output = self.outputs()
        self.assertIn("RESULT 42 3 [1] 7 5 alive b xml.dom False False False", output)
        notice = latest_user(self.provider.requests[-2]["request"])
        self.assertIn("re-run from source", notice)
        self.assertIn("helper", notice)

    def test_a_state_file_in_the_old_format_still_restores_its_variables(self):
        session = self.app.session()
        self.turn(session, "define")
        self.saved_after(session, 0)
        legacy = {"cwd": os.getcwd(), "names": {"legacy": pickle.dumps(5)}}
        self.state(session).write_bytes(pickle.dumps(legacy))
        self.turn(session, "legacy")
        self.assertIn("LEGACY 5", self.outputs())