"""Retrieval keeps exact counts, scoped pages, and Unicode text offsets through the daemon."""

import base64
import hashlib
import json
import re
import unittest

from harness import Albedo, Provider, Reply, operation_id, python, text
from image_limits_test import png


class RetrievalTests(unittest.TestCase):
    def setUp(self):
        self.code = None
        self.tool = None
        self.arguments = None

        def script(request):
            last = request["input"][-1]
            if last.get("type") == "function_call_output":
                return text("done")
            if "<newly-evicted-history>" in str(last.get("content", "")):
                return text("fold-only Café")
            if self.code is not None:
                return python(self.code)
            if self.tool is not None:
                return Reply(
                    "python", tool_name=self.tool, tool_arguments=self.arguments
                )
            return text("done")

        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)
        self.session = self.app.session()

    def turn(self, session, prompt):
        before = len(self.provider.requests)
        self.app.prompt(session, prompt).close()
        self.app.idle(session)
        return [entry["request"] for entry in self.provider.requests[before:]]

    def retrieve(self, name, arguments):
        self.tool, self.arguments = name, arguments
        requests = self.turn(self.session, "retrieve fixture")
        output = next(
            item["output"]
            for item in reversed(requests[-1]["input"])
            if item.get("type") == "function_call_output"
        )
        return json.loads(output)

    def test_transcript_pages_count_matches_and_preserve_grapheme_offsets(self):
        target = self.app.session()
        for prompt in ["é👩‍💻tail Café\r", "second Café", "third Café"]:
            self.turn(target, prompt)
        self.code = f"""
agent = await agents.get({target!r})
found = await agent.search_messages('CAFÉ', limit=1)
assert found.count == 3, found
assert len(found.rows) == 1 and found.next_offset == 1, found
seq = found.rows[0]['seq']
second = await agent.search_messages('CAFÉ', offset=1, limit=1)
assert second.count == 3 and second.rows[0]['seq'] > seq, second
end = await agent.search_messages('CAFÉ', offset=3, limit=1)
assert end.count == 3 and end.rows == [] and end.next_offset is None, end
prefix = len(f'[row #{{seq}}]\\n[user]\\n')
accent = await agent.messages(seq=seq, offset=prefix, limit=1)
assert accent.content == 'é', accent
emoji = await agent.messages(seq=seq, offset=prefix+1, limit=1)
assert emoji.content == '👩‍💻', emoji
boundary = await agent.messages(seq=seq, offset=prefix+11, limit=1)
assert boundary.content == '\\r\\n', boundary
separator = await agent.messages(seq=seq, offset=prefix+12, limit=1)
assert separator.content == '\\n', separator
beyond = await agent.messages(seq=seq, offset=100000, limit=1)
assert beyond.content == '' and beyond.next_offset is None, beyond
try:
    await agent.messages(seq=seq, offset=-1)
    assert False, 'negative offset was accepted'
except AgentsError as error:
    assert 'nonnegative' in str(error), error
print('RETRIEVAL_PAGES_OK')
"""
        requests = self.turn(self.session, "verify source retrieval")
        output = next(
            item["output"]
            for item in reversed(requests[-1]["input"])
            if item.get("type") == "function_call_output"
        )
        self.assertIn("RETRIEVAL_PAGES_OK", output)

    def test_a_repeated_image_reads_back_with_the_same_identity(self):
        target = self.app.session()
        same, other = (
            base64.b64encode(png(width, height)).decode()
            for width, height in ((2, 3), (4, 5))
        )
        for words, data in (("first", same), ("again", same), ("new", other)):
            self.app.api(
                f"/sessions/{target}/inputs/{operation_id()}",
                {
                    "kind": "message",
                    "text": words,
                    "images": [{"mime_type": "image/png", "data": data}],
                },
                method="PUT",
            ).close()
            self.app.idle(target)
        self.code = f"print((await (await agents.get({target!r})).messages(limit=8000)).content)"
        requests = self.turn(self.session, "read the target")
        output = next(
            item["output"]
            for item in reversed(requests[-1]["input"])
            if item.get("type") == "function_call_output"
        )
        labels = re.findall(
            r"\[user with image/png \d+x\d+ #([0-9a-f]{8})\]",
            json.loads(output)["output"],
        )
        expected = hashlib.sha256(same.encode()).hexdigest()[:8]
        self.assertEqual(labels[:2], [expected, expected], output)
        self.assertEqual(len(labels), 3, output)
        self.assertNotEqual(labels[2], expected)

    def test_assistant_rows_read_as_text_and_reads_and_searches_take_bounds(self):
        # The code has quotes and a newline, which a provider's JSON escapes.
        self.code = 'print("ROWMARK")\nprint(2)'
        self.turn(self.session, "run the marker")
        self.code = None
        pattern = 'print("ROWMARK")'
        found = self.retrieve("transcript_grep", {"pattern": pattern})
        # The search's own call row also holds the pattern; it is not counted.
        self.assertEqual(found["count"], 1, found)
        called = [row for row in found["rows"] if "[call python]" in row["preview"]]
        self.assertEqual(len(called), 1, found)
        seq = called[0]["seq"]
        self.assertEqual(called[0]["kind"], "assistant")
        self.assertNotIn("\\", called[0]["preview"])
        # through stops after the row; the text is the call, not provider JSON.
        row = self.retrieve(
            "transcript_read", {"seq": seq, "through": seq, "limit": 8000}
        )
        self.assertEqual(
            row["content"],
            f'[row #{seq}]\n[assistant]\n[call python]\ncode:\nprint("ROWMARK")\n'
            "print(2)\ntimeout_ms: 60000",
        )
        self.assertIsNone(row["next_offset"])
        # Without through the page runs on into the rows after it.
        onward = self.retrieve("transcript_read", {"seq": seq, "limit": 8000})
        self.assertGreater(onward["content"].count("[row #"), 1, onward)
        after = self.retrieve(
            "transcript_grep", {"pattern": "ROWMARK", "after_seq": seq}
        )
        self.assertTrue(after["rows"])
        self.assertTrue(all(item["seq"] > seq for item in after["rows"]), after)
        before = self.retrieve(
            "transcript_grep", {"pattern": "ROWMARK", "before_seq": seq}
        )
        self.assertTrue(all(item["seq"] < seq for item in before["rows"]), before)
        within = self.retrieve(
            "transcript_grep",
            {"pattern": "ROWMARK", "after_seq": seq - 1, "before_seq": seq + 1},
        )
        self.assertEqual([item["seq"] for item in within["rows"]], [seq], within)

    def test_python_reads_and_searches_take_the_same_bounds(self):
        target = self.app.session()
        for prompt in ["first marker", "second marker", "third marker"]:
            self.turn(target, prompt)
        self.code = f"""
mine = await agents.self.search_messages('SELFECHO')
assert mine.count == 0, mine
agent = await agents.get({target!r})
found = await agent.search_messages('marker')
users = [row['seq'] for row in found.rows if row['kind'] == 'user']
assert len(users) == 3, found
first, second, third = users
after = await agent.search_messages('marker', after_seq=first)
assert [r['seq'] for r in after.rows if r['kind'] == 'user'] == [second, third], after
before = await agent.search_messages('marker', before_seq=third)
assert [r['seq'] for r in before.rows if r['kind'] == 'user'] == [first, second], before
one = await agent.messages(first, through=first)
assert one.content.startswith(f'[row #{{first}}]'), one
assert f'[row #{{second}}]' not in one.content and one.next_offset is None, one
try:
    await agent.messages(second, through=first)
    assert False, 'through before seq was accepted'
except AgentsError as error:
    assert 'through' in str(error), error
print('RETRIEVAL_BOUNDS_OK')
"""
        requests = self.turn(self.session, "verify the bounds")
        output = next(
            item["output"]
            for item in reversed(requests[-1]["input"])
            if item.get("type") == "function_call_output"
        )
        self.assertIn("RETRIEVAL_BOUNDS_OK", output)

    def test_lcm_search_pages_sources_and_nodes_independently_inside_selected_range(
        self,
    ):
        for prompt in ["first Café", "second Café", "third Café", "fourth Café"]:
            self.turn(self.session, prompt)
        with self.app.api(
            f"/sessions/{self.session}/compaction",
            {"strategy": "lcm"},
        ) as response:
            result = json.load(response)
        self.assertEqual(result["state"], "compacted")
        self.app.idle(self.session)
        listed = self.retrieve("lcm_list", {"limit": 1})
        self.assertGreater(listed["total"], 0)
        self.assertEqual(len(listed["folds"]), 1)
        node = listed["folds"][0]
        self.assertTrue(node["frontier"])
        scoped = self.retrieve(
            "lcm_grep", {"pattern": "CAFÉ", "summary_id": node["id"], "limit": 1}
        )
        self.assertGreater(scoped["source_count"], 1)
        self.assertEqual(scoped["node_count"], 1)
        self.assertEqual(len(scoped["sources"]), 1)
        self.assertEqual(len(scoped["nodes"]), 1)
        self.assertEqual(scoped["next_offset"], 1)
        source = scoped["sources"][0]
        self.assertEqual(source["node_id"], node["id"])
        self.assertGreaterEqual(source["seq"], node["first_seq"])
        self.assertLessEqual(source["seq"], node["last_seq"])
        second = self.retrieve(
            "lcm_grep",
            {"pattern": "CAFÉ", "summary_id": node["id"], "limit": 1, "offset": 1},
        )
        self.assertEqual(second["source_count"], scoped["source_count"])
        self.assertEqual(second["node_count"], 1)
        self.assertEqual(len(second["sources"]), 1)
        self.assertEqual(second["nodes"], [])
        self.assertEqual(second["next_offset"], 2)
        summary_only = self.retrieve(
            "lcm_grep", {"pattern": "fold-only", "summary_id": node["id"], "limit": 1}
        )
        self.assertEqual(summary_only["source_count"], 0)
        self.assertEqual(summary_only["node_count"], 1)
        self.assertEqual(summary_only["sources"], [])
        self.assertEqual(len(summary_only["nodes"]), 1)
        self.assertIsNone(summary_only["next_offset"])
        page = self.retrieve("lcm_expand", {"id": node["id"], "limit": 1})
        self.assertEqual(page["content"], "[")
        self.assertEqual(page["next_offset"], 1)
        continued = self.retrieve(
            "lcm_expand", {"id": node["id"], "limit": 1, "offset": 1}
        )
        self.assertEqual(continued["content"], "s")
        outside = self.retrieve(
            "lcm_grep", {"pattern": "CAFÉ", "summary_id": node["id"], "offset": 1000}
        )
        self.assertEqual(outside["source_count"], scoped["source_count"])
        self.assertEqual(outside["sources"], [])
        self.assertIsNone(outside["next_offset"])