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

import json
import unittest
import urllib.parse

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)
        entry = [
            entry
            for entry in self.app.history(self.session)["items"]
            if entry["kind"] == "tool_result" and entry["tool"]["name"] == "python"
        ][-1]

        def referenced_json(field):
            [reference] = [
                part["reference"]
                for part in entry["content"]
                if part["kind"] == "reference" and part["reference"]["field"] == field
            ]
            pieces, offset, token = [], 0, None
            for _ in range(100):
                path = reference["url"]
                if token is not None:
                    path += "?" + urllib.parse.urlencode({"next": token})
                with self.app.api(path) as response:
                    page = json.load(response)
                for chunk in page["parts"]:
                    if chunk["field"] != field:
                        continue
                    self.assertEqual(chunk["offset_bytes"], offset)
                    self.assertEqual(chunk["encoding"], "utf8")
                    pieces.append(chunk["text"])
                    offset += len(chunk["text"].encode())
                token = page["next"]
                if token is None:
                    break
            else:
                self.fail(f"{field} content pagination did not terminate")
            self.assertEqual(offset, reference["bytes"])
            return json.loads("".join(pieces))

        results = [
            json.loads(part["value"])
            for part in entry["content"]
            if part["kind"] == "json" and part["field"] == "result"
        ]
        result = results[0] if results else json.loads(referenced_json("result"))
        traces = [part["trace"] for part in entry["content"] if part["kind"] == "trace"]
        trace = traces[0] if traces else referenced_json("trace")
        self.assertEqual(result["status"], status, result)
        return result, 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)