import json
from types import SimpleNamespace

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage

from artemis.llm.prefix_cache import PromptPrefixTracker
from artemis.llm.router import ModelEndpoint
from artemis.services.llm import RobustChatModelWrapper


def test_prefix_diagnostics_distinguish_append_rewrite_and_tool_changes_without_payloads():
    tracker = PromptPrefixTracker()
    messages = [SystemMessage(content="private rules"), HumanMessage(content="private screenshot")]
    options = {"tools": [{"name": "click"}, {"name": "observe"}]}
    assert tracker.observe("task:flash", messages, options)["cold_start"]
    messages.extend(
        [
            AIMessage(content="", tool_calls=[{"name": "click", "args": {"x": 5}, "id": "c1"}]),
            ToolMessage(content="clicked", tool_call_id="c1"),
        ]
    )
    extended = tracker.observe("task:flash", messages, options)
    assert extended["append_only"] and extended["shared_prefix_messages"] == 2
    messages[1] = HumanMessage(content="compressed summary")
    rewritten = tracker.observe("task:flash", messages, options)
    assert rewritten["first_changed_message"] == 2 and not rewritten["append_only"]
    reordered = tracker.observe("task:flash", messages, {"tools": list(reversed(options["tools"]))})
    assert reordered["tools_or_format_changed"] and not reordered["append_only"]
    assert tracker.observe("task:checker", messages, options)["cold_start"]
    assert "private" not in json.dumps(tracker._previous)


def test_bound_wrappers_reuse_task_prefix_tracker_after_message_adaptation():
    class Model:
        kwargs = {"tools": [{"type": "function", "name": "observe"}]}

        def bind_tools(self, *_args, **_kwargs):
            return self

    ctx = SimpleNamespace(data_engine=SimpleNamespace(current_session_id="session"))
    wrapper = RobustChatModelWrapper(Model(), ctx=ctx, endpoint=ModelEndpoint(), role="flash")
    messages = [SystemMessage(content="rules"), HumanMessage(content="step 1")]

    def diagnostics(items):
        bound = wrapper.bind_tools([])  # Flash binds again for each call
        args, kwargs = bound._prepare_model_call((items,), {})
        return bound._diagnostics(args, kwargs)

    first = diagnostics(messages)
    messages.append(HumanMessage(content="step 2"))
    second = diagnostics(messages)
    assert first["prefix_cache"]["cold_start"]
    assert second["prefix_cache"]["append_only"]
    assert first["prompt_cache_key"] == second["prompt_cache_key"]
    assert not any("rules" in str(value) for value in ctx.prompt_prefix_tracker._previous.values())

    ctx.transcript_ledger = SimpleNamespace(
        context_phase=1, last_compaction={"reason": "context_pressure"}
    )
    third = diagnostics(messages[:1] + [HumanMessage(content="history summary")])
    assert third["prefix_cache"]["context_phase"] == 1
    assert third["prefix_cache"]["last_compaction"]["reason"] == "context_pressure"
    assert not third["prefix_cache"]["append_only"]