"""Automatic compaction must stay visible through helper calls and clear on run end."""

from datetime import datetime
import json
import queue
import threading
import unittest

from harness import Albedo, Provider, Reply, error, text, wait_until
from provider_events import responses_events


class CompactionProgressTests(unittest.TestCase):
    def setUp(self):
        self.armed = False
        self.reject_replay = False
        self.fail_summaries = False
        self.failures = []
        self.gates = queue.Queue()
        self.releases = []
        self.hold_main = False
        self.main_entered = threading.Event()
        self.main_release = threading.Event()
        self.provider = Provider(self.reply, chunk_size=64_000)
        self.addCleanup(self.provider.close)
        self.app = Albedo(self.provider, protocol="responses")
        self.app.__enter__()
        self.addCleanup(self.app.__exit__, None, None, None)
        self.addCleanup(self.release_all)
        self.session = self.app.session()
        self.configure({"context_window": 2_000_000})

        for prompt in ("history one", "history two", "history three"):
            self.turn(prompt)
        with self.app.api(
            f"/sessions/{self.session}/compaction", {"strategy": "rolling"}
        ) as response:
            self.assertEqual(json.load(response)["state"], "compacted")
        self.app.idle(self.session)

        # Summary rendering caps each assistant turn at 64k characters. Fourteen
        # separate turns exceed the 200k-token chunk budget even after LCM keeps
        # its newest large unit verbatim. The high window keeps rolling quiet.
        for index in range(14):
            self.turn(f"large history {index}")
        self.turn("small newest unit")
        self.configure(
            {
                "context_window": 100_000,
                "selection": {"extensions": {"coding-lcm": True}},
            }
        )
        self.reload()
        self.armed = True

    @staticmethod
    def helper(request):
        encoded = json.dumps(request)
        if (
            "Rewrite your compaction notes now" in encoded
            or "You are keeping notes for yourself across a compaction" in encoded
        ):
            return "notes", False
        if (
            "These are memories of consecutive parts" in encoded
            or "You are consolidating these ordered memory segments" in encoded
        ):
            return "merging", False
        if "Write the memory of one stretch of this conversation" in encoded:
            return "summarizing", True
        if "<newly-evicted-history>" in encoded:
            return "summarizing", False
        return None, False

    def reply(self, request):
        stage, replay = self.helper(request)
        if stage is not None:
            if not self.armed:
                return text("old rolling summary")
            if self.fail_summaries or (self.reject_replay and replay):
                self.failures.append(stage)
                return error(400, "summary refused")
            release, returned = threading.Event(), threading.Event()
            self.releases.append(release)
            self.gates.put((stage, release, returned))
            if not release.wait(30):
                raise AssertionError(f"{stage} helper was not released")
            response = text("compressed memory")

            def events():
                try:
                    yield from responses_events(response, "", "", 64_000)
                finally:
                    returned.set()

            return Reply("raw", events=events())
        if self.armed and self.hold_main:
            self.main_entered.set()
            if not self.main_release.wait(30):
                raise AssertionError("main assistant request was not released")
        prompt = request["input"][-1].get("content", "")
        if isinstance(prompt, str) and prompt.startswith("large history "):
            return text("A" * 64_000)
        if prompt == "small newest unit":
            return text(
                "ordinary answer", usage={"input_tokens": 250_000, "output_tokens": 2}
            )
        return text("ordinary answer")

    def release_all(self):
        self.armed = False
        self.hold_main = False
        for release in self.releases:
            release.set()
        self.main_release.set()

    def read(self):
        with self.app.api(f"/sessions/{self.session}?tail=0") as response:
            return json.load(response)

    def configure(self, fields):
        route = f"/sessions/{self.session}?view=configuration"
        with self.app.api(route) as response:
            response.read()
            etag = response.headers["ETag"]
        self.app.api(route, fields, method="PATCH", headers={"If-Match": etag}).close()

    def reload(self):
        self.app.api(f"/sessions/{self.session}/reload", {"target": "session"}).close()

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

    def gate(self, stage):
        try:
            actual, release, returned = self.gates.get(timeout=20)
        except queue.Empty:
            self.fail(f"no {stage} helper: {self.read()['status']}")
        self.assertEqual(actual, stage)
        self.addCleanup(release.set)
        return release, returned

    def progress(self, stage, completed, total, started=None, active=None):
        def matches():
            status = self.read()["status"]
            progress = status.get("compaction")
            return (
                status["phase"] == "compacting"
                and progress is not None
                and progress["stage"] == stage
                and progress["completed"] == completed
                and progress["total"] == total
                and (active is None or progress.get("active") == active)
            )

        wait_until(matches, 20, f"missing {stage} {completed}/{total} progress")
        status = self.read()["status"]
        progress = status["compaction"]
        self.assertEqual(progress["strategy"], "coding-lcm")
        self.assertIsNotNone(datetime.fromisoformat(progress["started_at"]).tzinfo)
        if started is not None:
            self.assertEqual(progress["started_at"], started)
        return status

    def restore_rolling(self):
        self.configure(
            {
                "context_window": 2_000_000,
                "selection": {"extensions": {"rolling": True}},
            }
        )
        self.reload()

    def test_automatic_switch_reports_chunks_merge_notes_then_main_generation(self):
        # Reject the cache-replay optimization so this turn exercises the real
        # source chunk path, including its concurrent summaries and final merge.
        self.reject_replay = True
        self.hold_main = True
        self.app.prompt(self.session, "compact during an ordinary turn").close()
        first, first_returned = self.gate("summarizing")
        second, _ = self.gate("summarizing")
        status = self.progress("summarizing", 0, 2, active=2)
        started, run_id = status["compaction"]["started_at"], status["run_id"]
        first.set()
        self.assertTrue(first_returned.wait(20))
        self.progress("summarizing", 1, 2, started, active=1)
        second.set()

        merge, _ = self.gate("merging")
        self.progress("merging", 0, 1, started, active=1)
        merge.set()
        notes, _ = self.gate("notes")
        total = self.read()["status"]["compaction"]["total"]
        self.assertIsInstance(total, int)
        self.assertGreaterEqual(total, 1)
        for completed in range(total):
            if completed:
                notes, _ = self.gate("notes")
            self.progress("notes", completed, total, started)
            notes.set()
        self.assertTrue(self.main_entered.wait(20), "main request was never reached")
        status = self.read()["status"]
        self.assertEqual(status["run_id"], run_id)
        self.assertEqual(status["phase"], "generating")
        self.assertIsNone(status.get("compaction"))
        self.main_release.set()
        final = self.app.idle(self.session)
        self.assertIsNone(final["status"].get("compaction"))

    def test_stop_clears_progress_and_late_summary_cannot_claim_the_next_run(self):
        self.app.prompt(self.session, "compact then stop").close()
        release, returned = self.gate("summarizing")
        second, second_returned = self.gate("summarizing")
        status = self.progress("summarizing", 0, 2)
        captured = self.read()
        self.app.api(
            f"/sessions/{self.session}/interrupt",
            {
                "run_id": status["run_id"],
                "through_input_order": captured["input_order"],
            },
        ).close()
        stopped = self.app.idle(self.session)
        self.assertIsNone(stopped["status"].get("compaction"))

        self.restore_rolling()
        self.hold_main = True
        self.app.prompt(self.session, "new run after Stop").close()
        self.assertTrue(self.main_entered.wait(20))
        current = self.read()["status"]
        self.assertNotEqual(current["run_id"], status["run_id"])
        self.assertEqual(current["phase"], "generating")
        self.assertIsNone(current.get("compaction"))
        release.set()
        second.set()
        self.assertTrue(
            returned.wait(20), "old summary did not finish its provider gate"
        )
        self.assertTrue(second_returned.wait(20))
        current = self.read()["status"]
        self.assertEqual(current["phase"], "generating")
        self.assertIsNone(current.get("compaction"))
        self.main_release.set()
        self.app.idle(self.session)

    def test_nonretryable_summary_failure_preserves_history_and_rolling_projection(
        self,
    ):
        history = self.app.history(self.session)["items"]
        self.fail_summaries = True
        self.app.prompt(self.session, "a summary provider failure").close()
        failed = self.app.idle(self.session)
        self.assertGreaterEqual(
            len(self.failures), 2, "the source summary batch must fail"
        )
        self.assertEqual(failed["status"]["phase"], "idle")
        self.assertIsNone(failed["status"].get("compaction"))
        after = self.app.history(self.session)["items"]
        self.assertTrue(
            after[: len(history)] == history, "failed summary changed source history"
        )
        self.restore_rolling()
        self.turn("request after failed compaction")
        self.assertIn(
            "old rolling summary",
            json.dumps(self.provider.requests[-1]["request"]["input"]),
        )