from __future__ import annotations

from dataclasses import asdict

import pytest

from agent.core.runner import AgentRunner, RunConfig
from agent.orchestration.models import CompletionPolicy, TaskRecord
from agent.tools.builtin.context import get_current_context
from agent.tools.builtin.knowledge import KnowledgeConfig
from agent.trace.goal_models import GoalTree
from agent.trace.models import Trace
from agent.trace.store import FileSystemTraceStore


def _explicit_config(**overrides):
    values = {
        "agent_type": "planner",
        "completion_policy": CompletionPolicy.EXPLICIT_VALIDATION,
        "name": "explicit mission",
    }
    values.update(overrides)
    return RunConfig(**values)


def test_legacy_task_goal_id_is_ignored_when_loading_ledger_record():
    record = TaskRecord.from_dict(
        {
            "task_id": "task",
            "goal_id": "obsolete-projection",
            "parent_task_id": None,
            "display_path": "0",
            "specs": [
                {
                    "version": 1,
                    "objective": "mission",
                    "acceptance_criteria": [
                        {"criterion_id": "done", "description": "complete"}
                    ],
                }
            ],
        }
    )

    assert "goal_id" not in record.__dict__
    assert "goal_id" not in asdict(record)


@pytest.mark.asyncio
async def test_new_traces_persist_only_the_selected_plan_model(tmp_path):
    store = FileSystemTraceStore(str(tmp_path))
    runner = AgentRunner(trace_store=store)

    explicit, explicit_tree, _ = await runner._prepare_new_trace(
        [],
        _explicit_config(new_trace_id="explicit"),
    )
    legacy, legacy_tree, _ = await runner._prepare_new_trace(
        [],
        RunConfig(new_trace_id="legacy", name="legacy mission"),
    )

    assert explicit.agent_role == "planner"
    assert explicit.parent_goal_id is None
    assert explicit_tree is None
    assert await store.get_goal_tree("explicit") is None
    assert legacy.agent_role == "legacy"
    assert legacy_tree is not None
    assert await store.get_goal_tree("legacy") is not None


@pytest.mark.asyncio
async def test_explicit_resume_ignores_historical_goal_file(tmp_path):
    store = FileSystemTraceStore(str(tmp_path))
    await store.create_trace(
        Trace(
            trace_id="explicit",
            mode="agent",
            task="mission",
            agent_type="planner",
            agent_role="planner",
        )
    )
    await store.update_goal_tree("explicit", GoalTree(mission="obsolete projection"))
    runner = AgentRunner(trace_store=store)

    _, loaded_tree, _ = await runner._prepare_existing_trace(
        _explicit_config(trace_id="explicit")
    )

    assert loaded_tree is None
    assert (await store.get_goal_tree("explicit")).mission == "obsolete projection"


@pytest.mark.asyncio
async def test_explicit_current_context_uses_task_ledger_read_model(tmp_path):
    class Coordinator:
        async def task_context(self, root_trace_id):
            assert root_trace_id == "root"
            return "## Current Task Plan\n\n- 0 [running] deliver result"

    store = FileSystemTraceStore(str(tmp_path))
    await store.create_trace(
        Trace(
            trace_id="planner",
            mode="agent",
            agent_role="planner",
            context={"root_trace_id": "root"},
        )
    )
    runner = AgentRunner(trace_store=store, task_coordinator=Coordinator())

    result = await get_current_context(
        {
            "runner": runner,
            "trace_id": "planner",
            "root_trace_id": "root",
            "completion_policy": CompletionPolicy.EXPLICIT_VALIDATION.value,
        }
    )

    assert "Current Task Plan" in result.output
    assert "deliver result" in result.output
    assert "GoalTree" not in result.output
    compression_context = await runner._current_execution_context(
        "planner",
        None,
        _explicit_config(),
    )
    assert compression_context == "## Current Task Plan\n\n- 0 [running] deliver result"


@pytest.mark.asyncio
async def test_explicit_overflow_skips_goal_level_one(monkeypatch, tmp_path):
    store = FileSystemTraceStore(str(tmp_path))
    await store.create_trace(
        Trace(trace_id="root", mode="agent", agent_role="planner")
    )
    runner = AgentRunner(trace_store=store)
    history = [{"role": "user", "content": "large context " * 20}]

    monkeypatch.setattr(
        "agent.core.runner.CompressionConfig.get_max_tokens",
        lambda _self, _model: 1,
    )
    monkeypatch.setattr(
        "agent.core.runner.compress_completed_goals",
        lambda *_args, **_kwargs: (_ for _ in ()).throw(
            AssertionError("explicit mode must not run Goal Level 1")
        ),
    )

    updated, _, _, needs_level_two = await runner._manage_context_usage(
        "root",
        history,
        GoalTree(mission="must be ignored"),
        _explicit_config(
            knowledge=KnowledgeConfig(
                enable_extraction=False,
                enable_completion_extraction=False,
                enable_injection=False,
            )
        ),
        sequence=1,
        head_seq=0,
    )

    assert updated == history
    assert needs_level_two is True