"""Test suite for agentloop.

Run it directly (``python3 -m unittest discover -s tests`` from the artifact
root) or through check.py, which also runs an independent tamper probe.

The tamper tests are the teeth: a verifier that cannot be shown to fail is not
evidence.  Every mutation below is one a real editor or a buggy writer could
make, and each is asserted to be caught AND to be caught at the right record.
"""

from __future__ import annotations

import json
import os
import sys
import tempfile
import unittest

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
if ROOT not in sys.path:
    sys.path.insert(0, ROOT)

from agentloop import (  # noqa: E402
    AgentLoop,
    ModelReply,
    StopReason,
    Tool,
    ToolCall,
    ToolRegistry,
    ToolResult,
    Transcript,
    TranscriptError,
    ValidationError,
    canonical_json,
    estimate_tokens,
    validate,
    verify_file,
    verify_records,
)
from agentloop.transcript import GENESIS, KIND_ASSISTANT, KIND_FINAL, KIND_TOOL, KIND_USER  # noqa: E402


def add(a, b):
    return a + b


def make_registry():
    return ToolRegistry(
        [
            Tool("add", "Add two integers.", {
                "type": "object",
                "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
                "required": ["a", "b"],
                "additionalProperties": False,
            }, add),
            Tool("shout", "Uppercase a string.", {
                "type": "object",
                "properties": {"text": {"type": "string", "minLength": 1}},
                "required": ["text"],
            }, lambda text: text.upper()),
        ]
    )


def write_tmp(records, raw_lines=None):
    fd, path = tempfile.mkstemp(suffix=".jsonl")
    os.close(fd)
    with open(path, "wb") as fh:
        if raw_lines is not None:
            fh.write(("\n".join(raw_lines) + "\n").encode("utf-8"))
        else:
            fh.write("".join(canonical_json(r) + "\n" for r in records).encode("utf-8"))
    return path


# --------------------------------------------------------------------------
# transcript: the chain
# --------------------------------------------------------------------------


class TestTranscriptChain(unittest.TestCase):
    def test_chain_links_each_record_to_its_predecessor(self):
        tr = Transcript()
        tr.append(KIND_USER, {"text": "hi"})
        tr.append(KIND_ASSISTANT, {"text": "yo"})
        tr.append(KIND_TOOL, {"name": "add", "ok": True, "value": 5})
        self.assertEqual([r["seq"] for r in tr.records], [0, 1, 2])
        self.assertEqual(tr.records[0]["prev"], GENESIS)
        for earlier, later in zip(tr.records, tr.records[1:]):
            self.assertEqual(later["prev"], earlier["hash"])

    def test_genesis_hash_is_the_all_zero_constant(self):
        tr = Transcript()
        tr.append(KIND_USER, {"text": "x"})
        self.assertEqual(tr.records[0]["prev"], "0" * 64)

    def test_bytes_are_deterministic_and_reloadable(self):
        def build():
            tr = Transcript()
            tr.append(KIND_USER, {"text": "same"})
            tr.append(KIND_ASSISTANT, {"text": "same", "tool_calls": []})
            return tr

        self.assertEqual(build().bytes(), build().bytes())
        path = write_tmp(build().records)
        try:
            self.assertEqual(Transcript(path).bytes(), build().bytes())
        finally:
            os.unlink(path)

    def test_save_then_load_round_trips(self):
        tr = Transcript()
        tr.append(KIND_USER, {"text": "round trip"})
        with tempfile.TemporaryDirectory() as d:
            path = os.path.join(d, "t.jsonl")
            tr.save(path)
            self.assertEqual(Transcript(path).records, tr.records)
            self.assertEqual(verify_file(path)[:2], (True, []))

    def test_rejects_unknown_kind_and_bad_payload(self):
        tr = Transcript()
        with self.assertRaises(TranscriptError):
            tr.append("system", {"text": "nope"})
        with self.assertRaises(TranscriptError):
            tr.append(KIND_USER, "not a dict")
        with self.assertRaises(TranscriptError):
            tr.append(KIND_USER, {"bad": object()})

    def test_load_raises_on_bad_json_with_the_line_number(self):
        path = write_tmp(None, raw_lines=['{"seq": 0}', "{not json"])
        try:
            with self.assertRaises(TranscriptError) as ctx:
                Transcript(path)
            self.assertIn("line 2", str(ctx.exception))
        finally:
            os.unlink(path)


class TestTamperDetection(unittest.TestCase):
    """The teeth.  Each mutation must be caught, and at the named record."""

    def setUp(self):
        tr = Transcript()
        tr.append(KIND_USER, {"text": "add 2 and 3"})
        tr.append(KIND_ASSISTANT, {"text": "calling add", "tool_calls": []})
        tr.append(KIND_TOOL, {"name": "add", "ok": True, "value": 5})
        tr.append(KIND_FINAL, {"stop_reason": "no_tool_calls", "steps": 1, "tokens": 4})
        self.records = tr.records

    def test_clean_chain_verifies(self):
        ok, violations = verify_records(self.records)
        self.assertTrue(ok, violations)
        self.assertEqual(violations, [])

    def test_payload_rewrite_is_caught_at_that_record(self):
        tampered = json.loads(json.dumps(self.records))
        tampered[2]["payload"]["value"] = 999
        ok, violations = verify_records(tampered)
        self.assertFalse(ok)
        self.assertEqual(violations[0]["seq"], 2)
        self.assertEqual(violations[0]["reason"], "hash_mismatch")

    def test_single_character_edit_in_payload_is_caught(self):
        tampered = json.loads(json.dumps(self.records))
        tampered[0]["payload"]["text"] = "add 2 and 4"
        ok, violations = verify_records(tampered)
        self.assertFalse(ok)
        self.assertEqual(violations[0]["seq"], 0)
        self.assertEqual(violations[0]["reason"], "hash_mismatch")

    def test_dropped_middle_record_is_caught_as_a_seq_gap(self):
        tampered = json.loads(json.dumps(self.records))
        del tampered[1]
        ok, violations = verify_records(tampered)
        self.assertFalse(ok)
        self.assertEqual(violations[0]["seq"], 1)
        self.assertEqual(violations[0]["reason"], "seq_gap")

    def test_swapped_records_break_the_chain_at_the_first_swapped_index(self):
        """A raw swap moves record 2's seq into slot 1, so the seq check fires
        first.  The chain breaks and names index 1 -- which is the fault."""
        tampered = json.loads(json.dumps(self.records))
        tampered[1], tampered[2] = tampered[2], tampered[1]
        ok, violations = verify_records(tampered)
        self.assertFalse(ok)
        self.assertEqual(violations[0]["seq"], 1)
        self.assertEqual(violations[0]["reason"], "seq_gap")

    def test_reorder_plus_renumber_is_caught_as_a_prev_mismatch(self):
        """The harder attack: swap two records AND renumber them so the seq
        check passes.  Only the prev link can catch this, and it does."""
        tampered = json.loads(json.dumps(self.records))
        tampered[1], tampered[2] = tampered[2], tampered[1]
        tampered[1]["seq"], tampered[2]["seq"] = 1, 2
        ok, violations = verify_records(tampered)
        self.assertFalse(ok)
        self.assertEqual(violations[0]["seq"], 1)
        self.assertEqual(violations[0]["reason"], "prev_mismatch")

    def test_forged_record_with_a_recomputed_hash_still_breaks_the_chain(self):
        """An attacker who rewrites a payload AND recomputes that record's hash
        is still caught, because the NEXT record still points at the old hash."""
        from agentloop.transcript import _record_hash

        tampered = json.loads(json.dumps(self.records))
        tampered[2]["payload"]["value"] = 999
        tampered[2]["hash"] = _record_hash(
            tampered[2]["seq"], tampered[2]["kind"], tampered[2]["payload"], tampered[2]["prev"]
        )
        ok, violations = verify_records(tampered)
        self.assertFalse(ok)
        self.assertEqual(violations[0]["seq"], 3)
        self.assertEqual(violations[0]["reason"], "prev_mismatch")

    def test_hash_field_replaced_by_a_random_hex_is_caught(self):
        tampered = json.loads(json.dumps(self.records))
        tampered[1]["hash"] = "f" * 64
        ok, violations = verify_records(tampered)
        self.assertFalse(ok)
        self.assertEqual(violations[0]["seq"], 1)
        self.assertEqual(violations[0]["reason"], "hash_mismatch")

    def test_missing_field_is_reported_by_name(self):
        tampered = json.loads(json.dumps(self.records))
        del tampered[0]["prev"]
        ok, violations = verify_records(tampered)
        self.assertFalse(ok)
        self.assertEqual(violations[0]["reason"], "missing_fields")
        self.assertEqual(violations[0]["detail"], ["prev"])

    def test_non_object_record_is_caught(self):
        tampered = json.loads(json.dumps(self.records))
        tampered[1] = "i am not a record"
        ok, violations = verify_records(tampered)
        self.assertFalse(ok)
        self.assertEqual(violations[0]["reason"], "not_an_object")

    def test_bad_json_line_is_reported_with_its_line_number(self):
        path = write_tmp(None, raw_lines=[canonical_json(self.records[0]), "{oops"])
        try:
            ok, violations, count = verify_file(path)
            self.assertFalse(ok)
            self.assertEqual(violations[0]["reason"], "bad_json")
            self.assertEqual(violations[0]["line"], 2)
            self.assertEqual(count, 1)
        finally:
            os.unlink(path)

    def test_truncation_is_detected_only_against_an_external_count(self):
        """The honest limit: the chain alone cannot see a missing tail."""
        full = write_tmp(self.records)
        cut = write_tmp(self.records[:2])
        try:
            ok_full, _, n_full = verify_file(full)
            ok_cut, _, n_cut = verify_file(cut)
            self.assertTrue(ok_full)
            self.assertTrue(ok_cut)  # chain intact -- this is the documented limit
            self.assertLess(n_cut, n_full)
        finally:
            os.unlink(full)
            os.unlink(cut)

    def test_verify_file_reports_an_unreadable_path(self):
        ok, violations, count = verify_file("/nonexistent/nope.jsonl")
        self.assertFalse(ok)
        self.assertEqual(violations[0]["reason"], "unreadable")
        self.assertEqual(count, 0)


# --------------------------------------------------------------------------
# tool arguments
# --------------------------------------------------------------------------


class TestArgumentValidation(unittest.TestCase):
    def test_valid_arguments_pass(self):
        validate({"a": 2, "b": 3}, {
            "type": "object",
            "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
            "required": ["a", "b"],
        })

    def test_missing_required_property(self):
        with self.assertRaises(ValidationError) as ctx:
            validate({"a": 2}, {"type": "object", "properties": {"a": {"type": "integer"}},
                                 "required": ["a", "b"]})
        self.assertIn("b", str(ctx.exception))

    def test_wrong_type(self):
        with self.assertRaises(ValidationError):
            validate({"a": "2", "b": 3}, {
                "type": "object",
                "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
            })

    def test_bool_is_not_an_integer(self):
        """bool subclasses int; without the explicit guard True == 1."""
        with self.assertRaises(ValidationError):
            validate({"a": True, "b": 3}, {
                "type": "object",
                "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
            })

    def test_float_is_a_number_but_not_an_integer(self):
        validate(1.5, {"type": "number"})
        with self.assertRaises(ValidationError):
            validate(1.5, {"type": "integer"})

    def test_additional_properties_false_rejects_extras(self):
        with self.assertRaises(ValidationError) as ctx:
            validate({"a": 1, "b": 2, "c": 3}, {
                "type": "object",
                "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
                "additionalProperties": False,
            })
        self.assertIn("c", str(ctx.exception))

    def test_enum(self):
        validate("red", {"enum": ["red", "green"]})
        with self.assertRaises(ValidationError):
            validate("blue", {"enum": ["red", "green"]})

    def test_array_items_are_validated_recursively(self):
        schema = {"type": "array", "items": {"type": "integer"}, "minItems": 1}
        validate([1, 2, 3], schema)
        with self.assertRaises(ValidationError) as ctx:
            validate([1, "two"], schema)
        self.assertIn("[1]", str(ctx.exception))

    def test_min_items_and_max_items(self):
        with self.assertRaises(ValidationError):
            validate([], {"type": "array", "minItems": 1})
        with self.assertRaises(ValidationError):
            validate([1, 2], {"type": "array", "maxItems": 1})

    def test_string_bounds_and_pattern(self):
        with self.assertRaises(ValidationError):
            validate("", {"type": "string", "minLength": 1})
        with self.assertRaises(ValidationError):
            validate("toolong", {"type": "string", "maxLength": 3})
        with self.assertRaises(ValidationError):
            validate("abc", {"type": "string", "pattern": "^[0-9]+$"})
        validate("123", {"type": "string", "pattern": "^[0-9]+$"})

    def test_numeric_bounds(self):
        validate(5, {"type": "integer", "minimum": 5, "maximum": 5})
        with self.assertRaises(ValidationError):
            validate(4, {"type": "integer", "minimum": 5})
        with self.assertRaises(ValidationError):
            validate(6, {"type": "integer", "maximum": 5})

    def test_nested_object_properties_are_checked(self):
        schema = {
            "type": "object",
            "properties": {"outer": {"type": "object", "properties": {"n": {"type": "integer"}},
                                       "required": ["n"]}},
        }
        validate({"outer": {"n": 1}}, schema)
        with self.assertRaises(ValidationError) as ctx:
            validate({"outer": {"n": "x"}}, schema)
        self.assertIn("outer.n", str(ctx.exception))

    def test_unknown_type_in_schema_is_a_loud_error(self):
        with self.assertRaises(ValidationError):
            validate(1, {"type": "flobble"})

    def test_tool_call_validates_before_running(self):
        tool = Tool("add", "add", {
            "type": "object",
            "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
            "required": ["a", "b"],
        }, add)
        self.assertEqual(tool({"a": 2, "b": 3}), 5)
        with self.assertRaises(ValidationError):
            tool({"a": 2})

    def test_tool_wraps_an_unexpected_exception_in_toolerror(self):
        def boom():
            raise RuntimeError("kaboom")

        tool = Tool("boom", "booms", {"type": "object", "properties": {}}, boom)
        with self.assertRaises(Exception) as ctx:
            tool({})
        self.assertIn("kaboom", str(ctx.exception))

    def test_registry_rejects_duplicate_names(self):
        reg = ToolRegistry([Tool("add", "", {"type": "object", "properties": {}}, add)])
        with self.assertRaises(ValueError):
            reg.add(Tool("add", "", {"type": "object", "properties": {}}, add))

    def test_registry_manifest_is_sorted_and_complete(self):
        reg = make_registry()
        names = [m["name"] for m in reg.manifest()]
        self.assertEqual(names, ["add", "shout"])
        for m in reg.manifest():
            self.assertEqual(sorted(m), ["description", "name", "parameters"])


# --------------------------------------------------------------------------
# the loop
# --------------------------------------------------------------------------


class TestLoop(unittest.TestCase):
    def test_completes_when_the_model_stops_asking_for_tools(self):
        def model(messages, manifest):
            return ModelReply("done", [])

        result = AgentLoop(make_registry()).run("hi", model)
        self.assertEqual(result.stop_reason, StopReason.NO_TOOL_CALLS)
        self.assertTrue(result.ok)
        self.assertEqual(result.steps, 1)

    def test_executes_tool_calls_in_order_and_records_them(self):
        calls = [ToolCall("add", {"a": 2, "b": 3}), ToolCall("shout", {"text": "hi"})]
        seen = []

        def model(messages, manifest):
            if not seen:
                seen.append(1)
                return ModelReply("working", calls)
            return ModelReply("2+3=5 and HI", [])

        result = AgentLoop(make_registry()).run("hi", model)
        results = result.tool_results()
        self.assertEqual([r.name for r in results], ["add", "shout"])
        self.assertEqual([r.value for r in results], [5, "HI"])
        self.assertTrue(all(r.ok for r in results))

    def test_unknown_tool_is_surfaced_as_an_error_and_the_loop_continues(self):
        def model(messages, manifest):
            if len(messages) == 1:
                return ModelReply("try", [ToolCall("nope", {})])
            return ModelReply("recovered", [])

        result = AgentLoop(make_registry()).run("hi", model)
        self.assertEqual(result.stop_reason, StopReason.NO_TOOL_CALLS)
        (only,) = result.tool_results()
        self.assertFalse(only.ok)
        self.assertIn("no tool named 'nope'", only.error)

    def test_invalid_arguments_are_surfaced_as_an_error(self):
        def model(messages, manifest):
            if len(messages) == 1:
                return ModelReply("try", [ToolCall("add", {"a": 1})])
            return ModelReply("ok", [])

        result = AgentLoop(make_registry()).run("hi", model)
        (only,) = result.tool_results()
        self.assertFalse(only.ok)
        self.assertIn("invalid arguments", only.error)
        self.assertIn("b", only.error)

    def test_a_tool_that_raises_does_not_kill_the_loop(self):
        reg = ToolRegistry([Tool("boom", "booms", {"type": "object", "properties": {}},
                                 lambda: (_ for _ in ()).throw(RuntimeError("kaboom")))])

        def model(messages, manifest):
            if len(messages) == 1:
                return ModelReply("go", [ToolCall("boom", {})])
            return ModelReply("survived", [])

        result = AgentLoop(reg).run("hi", model)
        self.assertEqual(result.stop_reason, StopReason.NO_TOOL_CALLS)
        self.assertFalse(result.tool_results()[0].ok)

    def test_non_serialisable_return_value_is_refused(self):
        reg = ToolRegistry([Tool("obj", "returns an object", {"type": "object", "properties": {}},
                                 lambda: object())])

        def model(messages, manifest):
            if len(messages) == 1:
                return ModelReply("go", [ToolCall("obj", {})])
            return ModelReply("done", [])

        result = AgentLoop(reg).run("hi", model)
        (only,) = result.tool_results()
        self.assertFalse(only.ok)
        self.assertIn("non-serialisable", only.error)
        # and the transcript is still writable
        self.assertTrue(result.transcript.bytes())

    def test_step_budget_caps_the_loop_and_matches_the_record_count(self):
        def runaway(messages, manifest):
            return ModelReply("more", [ToolCall("shout", {"text": "x"})])

        result = AgentLoop(make_registry(), max_steps=4).run("hi", runaway)
        self.assertEqual(result.stop_reason, StopReason.STEP_BUDGET)
        self.assertEqual(result.steps, 4)
        self.assertFalse(result.ok)
        assistant_records = [r for r in result.transcript.records if r["kind"] == KIND_ASSISTANT]
        self.assertEqual(len(assistant_records), 4)
        self.assertEqual(result.steps, len(assistant_records))

    def test_token_budget_stops_the_loop_before_running_tools(self):
        def chatty(messages, manifest):
            return ModelReply("x" * 400, [ToolCall("shout", {"text": "x"})])

        result = AgentLoop(make_registry(), max_steps=50, max_tokens=20).run("hi", chatty)
        self.assertEqual(result.stop_reason, StopReason.TOKEN_BUDGET)
        self.assertEqual(result.tool_results(), [])
        self.assertGreater(result.tokens, 20)

    def test_sentinel_stops_the_loop(self):
        def model(messages, manifest):
            if len(messages) == 1:
                return ModelReply("first", [ToolCall("add", {"a": 1, "b": 1})])
            return ModelReply("the answer is 2 -- DONE", [])

        result = AgentLoop(make_registry(), stop_sentinel="DONE").run("hi", model)
        self.assertEqual(result.stop_reason, StopReason.SENTINEL)
        self.assertTrue(result.ok)

    def test_sentinel_wins_over_a_tool_call_in_the_same_turn(self):
        def model(messages, manifest):
            if len(messages) == 1:
                return ModelReply("first", [ToolCall("add", {"a": 1, "b": 1})])
            return ModelReply("done DONE", [ToolCall("add", {"a": 9, "b": 9})])

        result = AgentLoop(make_registry(), stop_sentinel="DONE").run("hi", model)
        self.assertEqual(result.stop_reason, StopReason.SENTINEL)
        self.assertEqual(len(result.tool_results()), 1)

    def test_model_exception_becomes_a_model_error_stop(self):
        def model(messages, manifest):
            raise RuntimeError("the model is down")

        result = AgentLoop(make_registry()).run("hi", model)
        self.assertEqual(result.stop_reason, StopReason.MODEL_ERROR)
        self.assertIn("the model is down", result.model_error)
        self.assertFalse(result.ok)

    def test_model_returning_the_wrong_type_becomes_a_model_error(self):
        result = AgentLoop(make_registry()).run("hi", lambda m, f: "not a reply")
        self.assertEqual(result.stop_reason, StopReason.MODEL_ERROR)

    def test_a_model_that_raises_mid_run_still_leaves_a_valid_transcript(self):
        state = {"n": 0}

        def model(messages, manifest):
            state["n"] += 1
            if state["n"] == 1:
                return ModelReply("go", [ToolCall("add", {"a": 1, "b": 2})])
            raise RuntimeError("died on turn 2")

        result = AgentLoop(make_registry()).run("hi", model)
        self.assertEqual(result.stop_reason, StopReason.MODEL_ERROR)
        self.assertTrue(verify_records(result.transcript.records)[0])
        self.assertEqual(result.transcript.records[-1]["kind"], KIND_FINAL)

    def test_transcript_record_order_is_user_assistant_tool_final(self):
        def model(messages, manifest):
            if len(messages) == 1:
                return ModelReply("go", [ToolCall("add", {"a": 1, "b": 1})])
            return ModelReply("done", [])

        result = AgentLoop(make_registry()).run("hi", model)
        kinds = [r["kind"] for r in result.transcript.records]
        self.assertEqual(kinds, [KIND_USER, KIND_ASSISTANT, KIND_TOOL, KIND_ASSISTANT, KIND_FINAL])

    def test_identical_runs_produce_byte_identical_transcripts(self):
        def model(messages, manifest):
            if len(messages) == 1:
                return ModelReply("go", [ToolCall("add", {"a": 40, "b": 2})])
            return ModelReply("42", [])

        a = AgentLoop(make_registry()).run("hi", model)
        b = AgentLoop(make_registry()).run("hi", model)
        self.assertEqual(a.transcript.bytes(), b.transcript.bytes())

    def test_a_supplied_transcript_is_appended_to_not_replaced(self):
        existing = Transcript()
        existing.append(KIND_USER, {"text": "an earlier session"})

        result = AgentLoop(make_registry()).run("hi", lambda m, f: ModelReply("done", []),
                                             transcript=existing)
        # the supplied record, then user + assistant + final = 4
        self.assertEqual(len(result.transcript), 4)
        self.assertEqual(result.transcript.records[0]["payload"]["text"], "an earlier session")
        self.assertEqual([r["kind"] for r in result.transcript.records][0], KIND_USER)
        self.assertTrue(verify_records(result.transcript.records)[0])

    def test_budgets_must_be_positive_integers(self):
        with self.assertRaises(ValueError):
            AgentLoop(make_registry(), max_steps=0)
        with self.assertRaises(ValueError):
            AgentLoop(make_registry(), max_steps=True)
        with self.assertRaises(ValueError):
            AgentLoop(make_registry(), max_tokens=0)

    def test_registry_must_be_a_toolregistry(self):
        with self.assertRaises(TypeError):
            AgentLoop({"add": add})

    def test_estimate_tokens_is_a_documented_floor(self):
        self.assertEqual(estimate_tokens(""), 0)
        self.assertEqual(estimate_tokens("ab"), 1)   # floor 1, not 0
        self.assertEqual(estimate_tokens("abcd"), 1)
        self.assertEqual(estimate_tokens("abcde"), 1)
        self.assertEqual(estimate_tokens("a" * 40), 10)


if __name__ == "__main__":
    unittest.main(verbosity=2)