"""Finalized traces survive buffer release, nested evaluation, and late edits."""

import json
import unittest

from harness import Albedo, Provider, python, text


class TracesTests(unittest.TestCase):
    def setUp(self):
        self.code = ""
        self.provider = Provider(
            lambda request: (
                python(self.code)
                if request["messages"][-1].get("role") == "user"
                else text("done")
            )
        )
        self.addCleanup(self.provider.close)
        self.app = Albedo(self.provider)
        self.app.__enter__()
        self.addCleanup(self.app.__exit__, None, None, None)
        self.session = self.app.session()

    def run_cell(self, code, status="ok"):
        self.code = code
        self.app.prompt(self.session, "execute this cell").close()
        self.app.idle(self.session)
        event = [
            event
            for event in self.app.events(self.session)
            if event.get("type") == "tool" and event.get("name") == "python"
        ][-1]
        result = json.loads(event["result"])
        self.assertEqual(result["status"], status, result)
        return result, event["trace"]

    def read_trace(self, id):
        result, _ = self.run_cell(
            f"import json\nprint(json.dumps(await cells.trace({id!r})))"
        )
        return json.loads(result["output"])

    def test_top_level_and_nested_traces_preserve_changes_for_all_outcomes(self):
        for status, ending in (
            ("ok", ""),
            ("error", "raise ValueError('failed')"),
            ("interrupted", "raise __import__('asyncio').CancelledError()"),
        ):
            with self.subTest(status=status):
                (self.app.workspace / "note.txt").write_text("before\n")
                (self.app.workspace / "binary.dat").unlink(missing_ok=True)
                source = (
                    """
import os
open('temporary.txt', 'w').write('after\\n')
os.replace('temporary.txt', 'note.txt')
open('note.txt').read()
open('binary.dat', 'wb').write(b'\\0binary')
"""
                    + ending
                )
                original, trace = self.run_cell(source, status)
                self.assertEqual(trace, self.read_trace(original["cell_id"]))
                self.assert_changes(trace)
                # The exact saved source runs as a separate nested cell.
                (self.app.workspace / "note.txt").write_text("before\n")
                (self.app.workspace / "binary.dat").unlink()
                parent, parent_trace = self.run_cell(
                    f"try:\n    await cells.run({original['cell_id']!r}, allow_partial=True)\n"
                    "except BaseException:\n    pass\nprint(cells.last_id)"
                )
                child = parent["output"].strip().splitlines()[-1]
                nested = self.read_trace(child)
                self.assert_changes(nested)
                self.run_cell(
                    f"assert (await cells.info({child!r}))['status'] == {status!r}"
                )
                self.assertEqual(parent_trace["changes"], [])

    def assert_changes(self, trace):
        changes = {change["path"].split("/")[-1]: change for change in trace["changes"]}
        self.assertNotIn("temporary.txt", changes)
        self.assertIn("-before", changes["note.txt"]["diff"])
        self.assertIn("+after", changes["note.txt"]["diff"])
        self.assertEqual(changes["binary.dat"]["kind"], "unavailable")
        self.assertTrue(
            any(
                activity["kind"] == "read" and activity["target"].endswith("note.txt")
                for activity in trace["activities"]
            )
        )

    def test_many_edit_heavy_cells_release_buffers_and_keep_old_traces(self):
        oldest = None
        oldest_trace = None
        for index in range(18):
            result, trace = self.run_cell(
                f"for index in range(20):\n    open(f'edit-{{index}}.txt', 'w').write({str(index)!r})\n"
                "open('edit-0.txt').read()"
            )
            self.assertTrue(trace["truncated"])
            self.assertTrue(trace["changes"])
            if oldest is None:
                oldest, oldest_trace = result["cell_id"], trace
        result, _ = self.run_cell("""
import __main__ as kernel
completed = [capture for capture in kernel.ARCHIVES.values() if capture.trace.sealed]
assert completed
assert all(not capture.trace.before and not capture.trace.activities for capture in completed)
print('released')
""")
        self.assertIn("released", result["output"])
        self.assertEqual(self.read_trace(oldest), oldest_trace)

    def test_late_background_edits_execute_without_changing_finalized_trace(self):
        result, trace = self.run_cell("""
import asyncio
import __main__ as kernel
late_gate = asyncio.Event()
async def late_edit():
    await late_gate.wait()
    open('late.txt', 'w').write('late\\n')
    open('late.txt').read()
late_task = asyncio.create_task(late_edit())
open('early.txt', 'w').write('early\\n')
""")
        self.run_cell(f"""
late_gate.set()
await late_task
capture = kernel.ARCHIVES[{result["cell_id"]!r}]
assert capture.trace.sealed
assert not capture.trace.before and not capture.trace.activities
""")
        self.assertEqual((self.app.workspace / "late.txt").read_text(), "late\n")
        self.assertEqual(self.read_trace(result["cell_id"]), trace)