"""A Python crash injects exactly one reset notice without corrupting history; a
daemon crash leaves a kernel with live work alive, so its next cell is told the
kernel reattached and the work carried on."""

import json
import urllib.error
import unittest

from harness import (
    Albedo,
    Provider,
    exclusive,
    is_note,
    prompt_text,
    python,
    text,
    turn_notes,
)

NOTICE = (
    "<system-note>The python kernel got reset and all variables are lost</system-note>"
)
REATTACHED = (
    "<system-note>The python kernel reattached; its variables and jobs carried on"
)


# exclusive: crashes and restarts the daemon
@exclusive
class KernelResetTests(unittest.TestCase):
    def setUp(self):
        cells = {
            # A live job holds the kernel through the crash, as a user's would.
            "first prompt": "job = run('sleep', '60')\nanswer = 7",
            "after restart": "print(answer, job.exit_code)",
            "lose kernel": "await job.stop()\nimport os; os._exit(1)",
            "recover kernel": "print('kernel-ready')",
        }

        def script(request):
            last = request["input"][-1]
            if last.get("role") == "user" and not is_note(last):
                code = cells.get(prompt_text(request))
                if code is not None:
                    return python(code)
            return text("ok")

        self.provider = Provider(script)
        self.addCleanup(self.provider.close)
        self.app = Albedo(self.provider, protocol="responses")
        self.app.__enter__()
        self.addCleanup(self.app.__exit__, None, None, None)

    def turn(self, session, prompt):
        """The notes the model saw in this turn."""
        self.app.prompt(session, prompt).close()
        self.app.idle(session)
        return turn_notes(self.provider.requests[-1]["request"])

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

    def test_kernel_crash_emits_one_shot_reset_notice_and_daemon_crash_reattaches(self):
        app = self.app
        session = app.session()
        self.assertEqual(self.turn(session, "first prompt"), [])
        unused = app.session()
        app.restart(crash=True)
        self.assertEqual(self.turn(unused, "unused session's first prompt"), [])
        with self.assertRaises(urllib.error.HTTPError) as rejected:
            app.prompt(session, "   ").close()
        self.assertEqual(rejected.exception.code, 400)
        [reattached] = self.turn(session, "after restart")
        self.assertTrue(reattached.startswith(REATTACHED), reattached)
        self.assertIn("Jobs still running", reattached)
        self.assertIn("7 None", self.outputs()[-1])
        self.assertEqual(self.turn(session, "ordinary turn"), [])
        self.assertEqual(self.turn(session, "lose kernel"), [])
        self.assertEqual(self.turn(session, "recover kernel"), [NOTICE])
        self.assertIn("kernel-ready", self.outputs()[-1])
        self.assertEqual(self.turn(session, "still alive"), [])
        users = [
            part["text"]
            for entry in app.history(session)["items"]
            if entry["kind"] == "user"
            for part in entry["content"]
            if part["kind"] == "text"
        ]
        self.assertEqual(
            users,
            [
                "first prompt",
                "after restart",
                "ordinary turn",
                "lose kernel",
                "recover kernel",
                "still alive",
            ],
        )
        with app.api("/sessions") as response:
            info = next(
                item for item in json.load(response)["items"] if item["id"] == session
            )
        self.assertEqual(info["name"], "still alive")