"""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 sqlite3
import unittest

from harness import alive, 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"
    "boxed = Box(9)\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 = (
    "import __main__, pickle\n"
    "assert __main__.answer == 7 and not hasattr(__main__, 'ARCHIVES')\n"
    "assert pickle.loads(pickle.dumps(Box(11))).v == 11\n"
    "print('RESULT', helper(21), Box(3).v, boxed.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())"
)


DECORATED = """def add(amount):
    def decorate(fn):
        return lambda: fn() + amount
    return decorate

def multiply(factor):
    def decorate(fn):
        return lambda: fn() * factor
    return decorate

def add_async(fn):
    async def wrapped():
        return await fn() + 100
    return wrapped

def mark(cls):
    cls.value += 100
    return cls

noise = '''unrelated
@not_a_decorator
'''
@(
    add(100)
)
# Keep stacked decorators in their original order.
@multiply(
    10
)
def decorated():
    return 7

@add_async
async def decorated_async():
    return 8

@mark
class Decorated:
    value = 9

print('DECORATED', decorated(), await decorated_async(), Decorated.value)
"""
RECALL_DECORATED = (
    "print('DECORATED', decorated(), await decorated_async(), Decorated.value)"
)


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

        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",
            # A budget of 0 keeps the daemon pressed, whose release drops the
            # values too big to save; unpressed, the kernel would be kept.
            prepare=lambda app: app.env.update(
                ALBEDO_IDLE_SECONDS=str(IDLE_SECONDS), ALBEDO_KERNEL_BUDGET_MB="0"
            ),
        )
        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 released(self, session):
        """Wait for the pressed sweep to stop the session's kernel."""
        with sqlite3.connect(self.app.home / "albedo.sqlite") as database:
            row = database.execute(
                "SELECT pid FROM kernel_links WHERE session=?", (session,)
            ).fetchone()
        # No row: the release already forgot the kernel.
        if row:
            wait_until(lambda: not alive(row[0]), 15, "the kernel was not released")
        self.assertTrue(self.state(session).exists())

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

    @exclusive
    def test_helpers_classes_imports_and_failed_cells_survive_a_restart(self):
        session = self.app.session()
        self.turn(session, "define")
        self.released(session)
        self.turn(session, "fail")
        self.released(session)
        self.turn(session, "recall")
        output = self.outputs()
        self.assertIn("RESULT 42 3 9 [1] 7 5 alive b xml.dom False False False", output)
        # Every definition re-ran, so the restore is lossless and silent.
        self.assertEqual(latest_user(self.provider.requests[-2]["request"]), "recall")

    def test_decorated_definitions_keep_behavior_after_release(self):
        session = self.app.session()
        self.turn(session, "decorated")
        self.assertIn(
            "DECORATED 170 108 109",
            self.provider.requests[-1]["request"]["input"][-1]["output"],
        )
        # A second release also exercises registration of replayed definitions.
        for _ in range(2):
            self.released(session)
            self.turn(session, "recall decorated")
            self.assertIn(
                "DECORATED 170 108 109",
                self.provider.requests[-1]["request"]["input"][-1]["output"],
            )