from __future__ import annotations

import pytest

from agent.core.runner import AgentRunner, RunConfig
from agent.tools.builtin.knowledge import KnowledgeConfig
from agent.trace.models import Message, Trace
from agent.trace.store import FileSystemTraceStore


async def _add(store, *, trace_id, role, sequence, parent, content):
    message = Message.create(
        trace_id=trace_id,
        role=role,
        sequence=sequence,
        parent_sequence=parent,
        content=content,
    )
    await store.add_message(message)
    return message


@pytest.mark.asyncio
async def test_level2_compression_carries_active_system_policies_across_reload(
    tmp_path,
):
    store = FileSystemTraceStore(str(tmp_path))
    await store.create_trace(Trace(trace_id="trace", mode="agent"))
    initial = await _add(
        store,
        trace_id="trace",
        role="system",
        sequence=1,
        parent=None,
        content="initial framework policy",
    )
    first_user = await _add(
        store,
        trace_id="trace",
        role="user",
        sequence=2,
        parent=1,
        content="mission",
    )
    await _add(
        store,
        trace_id="trace",
        role="assistant",
        sequence=3,
        parent=2,
        content={"text": "old work"},
    )
    await _add(
        store,
        trace_id="trace",
        role="system",
        sequence=4,
        parent=3,
        content="phase two policy",
    )
    await _add(
        store,
        trace_id="trace",
        role="user",
        sequence=5,
        parent=4,
        content="continue",
    )
    await store.update_trace("trace", head_sequence=5)
    runner = AgentRunner(trace_store=store)

    summary, next_sequence, carried = await runner._persist_compressed_main_path(
        trace_id="trace",
        original_head_sequence=5,
        next_sequence=10,
        summary_content="summary one",
    )
    await store.update_trace("trace", head_sequence=summary.sequence)

    assert summary.sequence == 11
    assert next_sequence == 12
    in_memory = runner._rebuild_history_after_compression(
        [
            initial.to_llm_dict(),
            first_user.to_llm_dict(),
            {"role": "assistant", "content": "old work"},
            {"role": "system", "content": "phase two policy"},
        ],
        summary.to_llm_dict(),
        carried_system_messages=carried,
    )
    assert [item["role"] for item in in_memory] == [
        "system",
        "user",
        "system",
        "user",
    ]
    assert in_memory[2]["content"] == "phase two policy"

    reloaded = FileSystemTraceStore(str(tmp_path))
    trace = await reloaded.get_trace("trace")
    main_path = await reloaded.get_main_path_messages("trace", trace.head_sequence)
    assert [message.role for message in main_path] == [
        "system",
        "user",
        "system",
        "user",
    ]
    assert [message.content for message in main_path if message.role == "system"] == [
        "initial framework policy",
        "phase two policy",
    ]


@pytest.mark.asyncio
async def test_level2_compression_does_not_duplicate_systems_before_first_user(
    tmp_path,
):
    store = FileSystemTraceStore(str(tmp_path))
    await store.create_trace(Trace(trace_id="trace", mode="agent"))
    parent = None
    for sequence, role, content in (
        (1, "system", "initial"),
        (2, "system", "startup policy"),
        (3, "user", "mission"),
        (4, "system", "phase two policy"),
        (5, "user", "continue"),
    ):
        await _add(
            store,
            trace_id="trace",
            role=role,
            sequence=sequence,
            parent=parent,
            content=content,
        )
        parent = sequence

    runner = AgentRunner(trace_store=store)
    summary, _, carried = await runner._persist_compressed_main_path(
        trace_id="trace",
        original_head_sequence=5,
        next_sequence=10,
        summary_content="summary",
    )
    await store.update_trace("trace", head_sequence=summary.sequence)
    path = await store.get_main_path_messages("trace", summary.sequence)

    assert [item["content"] for item in carried] == ["phase two policy"]
    assert [message.content for message in path if message.role == "system"] == [
        "initial",
        "startup policy",
        "phase two policy",
    ]


@pytest.mark.asyncio
async def test_repeated_level2_compression_preserves_policy_order_without_growth(
    tmp_path,
):
    store = FileSystemTraceStore(str(tmp_path))
    await store.create_trace(Trace(trace_id="trace", mode="agent"))
    parent = None
    for sequence, role, content in (
        (1, "system", "initial"),
        (2, "user", "mission"),
        (3, "system", "phase one"),
        (4, "user", "work"),
    ):
        await _add(
            store,
            trace_id="trace",
            role=role,
            sequence=sequence,
            parent=parent,
            content=content,
        )
        parent = sequence
    runner = AgentRunner(trace_store=store)
    summary_one, _, _ = await runner._persist_compressed_main_path(
        trace_id="trace",
        original_head_sequence=4,
        next_sequence=10,
        summary_content="summary one",
    )
    await _add(
        store,
        trace_id="trace",
        role="system",
        sequence=12,
        parent=summary_one.sequence,
        content="phase two",
    )
    await _add(
        store,
        trace_id="trace",
        role="user",
        sequence=13,
        parent=12,
        content="more work",
    )

    summary_two, _, _ = await runner._persist_compressed_main_path(
        trace_id="trace",
        original_head_sequence=13,
        next_sequence=20,
        summary_content="summary two",
    )
    await store.update_trace("trace", head_sequence=summary_two.sequence)
    path = await store.get_main_path_messages("trace", summary_two.sequence)

    assert [message.content for message in path if message.role == "system"] == [
        "initial",
        "phase one",
        "phase two",
    ]
    assert [message.role for message in path] == [
        "system",
        "user",
        "system",
        "system",
        "user",
    ]


@pytest.mark.asyncio
async def test_real_runner_level2_branch_persists_policy_across_store_reload(tmp_path):
    calls = 0

    async def llm_call(**kwargs):
        nonlocal calls
        calls += 1
        messages = kwargs["messages"]
        if any(
            "[[SUMMARY]]" in str(message.get("content", "")) for message in messages
        ):
            return {
                "content": "[[SUMMARY]] durable compressed state",
                "tool_calls": None,
                "finish_reason": "stop",
            }
        return {"content": "done", "tool_calls": None, "finish_reason": "stop"}

    store = FileSystemTraceStore(str(tmp_path))
    runner = AgentRunner(trace_store=store, llm_call=llm_call)
    config = RunConfig(
        new_trace_id="real-compression-trace",
        name="compression test",
        max_iterations=4,
        side_branch_max_turns=1,
        force_side_branch=["compression"],
        goal_compression="none",
        knowledge=KnowledgeConfig(
            enable_extraction=False,
            enable_completion_extraction=False,
            enable_injection=False,
        ),
    )

    events = [
        event
        async for event in runner.run(
            [
                {"role": "system", "content": "initial framework policy"},
                {"role": "user", "content": "mission"},
                {"role": "system", "content": "phase two policy"},
            ],
            config,
        )
    ]

    assert calls >= 2
    assert events[-1].status == "completed"
    reloaded = FileSystemTraceStore(str(tmp_path))
    trace = await reloaded.get_trace("real-compression-trace")
    main_path = await reloaded.get_main_path_messages(
        trace.trace_id, trace.head_sequence
    )
    system_contents = [
        str(message.content) for message in main_path if message.role == "system"
    ]
    assert len(system_contents) == 2
    assert "initial framework policy" in system_contents[0]
    assert "phase two policy" in system_contents[1]
    assert any(
        message.role == "user" and "durable compressed state" in str(message.content)
        for message in main_path
    )


@pytest.mark.asyncio
async def test_compression_does_not_reactivate_policy_outside_current_main_path(
    tmp_path,
):
    store = FileSystemTraceStore(str(tmp_path))
    await store.create_trace(Trace(trace_id="trace", mode="agent"))
    await _add(
        store,
        trace_id="trace",
        role="system",
        sequence=1,
        parent=None,
        content="initial",
    )
    await _add(
        store,
        trace_id="trace",
        role="user",
        sequence=2,
        parent=1,
        content="mission",
    )
    await _add(
        store,
        trace_id="trace",
        role="system",
        sequence=3,
        parent=2,
        content="active phase policy",
    )
    side_policy = Message.create(
        trace_id="trace",
        role="system",
        sequence=4,
        parent_sequence=3,
        branch_type="compression",
        branch_id="side",
        content="side branch policy must stay inactive",
    )
    await store.add_message(side_policy)
    await _add(
        store,
        trace_id="trace",
        role="user",
        sequence=5,
        parent=3,
        content="main path continues",
    )
    runner = AgentRunner(trace_store=store)

    summary, _, _ = await runner._persist_compressed_main_path(
        trace_id="trace",
        original_head_sequence=5,
        next_sequence=10,
        summary_content="summary",
    )
    path = await store.get_main_path_messages("trace", summary.sequence)

    policies = [str(message.content) for message in path if message.role == "system"]
    assert policies == ["initial", "active phase policy"]