"""Tests for shell-level compaction orchestration.

Spec source: context-compaction.implement-agent-compaction-shell
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import pytest
from pydantic_ai.agent import AgentRunResult
from pydantic_ai.messages import ModelRequest, ModelResponse, TextPart, UserPromptPart
from returns.result import Failure, Success

from anima.core.conversation import CompactionConfig
from anima.core.execution.state import TokenUsage
from anima.core.schemas import InternalError
from anima.shell import agent_compaction


@dataclass
class _StubAgent:
    calls: list[dict[str, Any]]
    output: str = "summary text"

    async def run(self, user_prompt: str, **kwargs: Any) -> AgentRunResult[str]:
        self.calls.append({"user_prompt": user_prompt, "kwargs": kwargs})
        return AgentRunResult(self.output)


def test_should_compact_respects_auto_flag() -> None:
    """auto=False disables compaction even during overflow."""
    usage = TokenUsage(prompt_tokens=180_000, completion_tokens=0, last_request_prompt_tokens=180_000)
    config = CompactionConfig(auto=False, reserved=20_000)

    result = agent_compaction.should_compact(
        token_usage=usage,
        compaction_config=config,
        model_context_limit=200_000,
    )

    assert result is False


def test_should_compact_uses_overflow_threshold() -> None:
    """Compaction is triggered exactly at context_limit-reserved."""
    usage = TokenUsage(prompt_tokens=180_000, completion_tokens=0, last_request_prompt_tokens=180_000)
    config = CompactionConfig(auto=True, reserved=20_000)

    result = agent_compaction.should_compact(
        token_usage=usage,
        compaction_config=config,
        model_context_limit=200_000,
    )

    assert result is True


def test_should_compact_ignores_completion_tokens() -> None:
    """Compaction threshold uses prompt_tokens only, not total_tokens.

    The context window limit constrains input (prompt) size.  Completion
    tokens are generated outside the window and must not inflate the
    overflow check.  With prompt_tokens=100k and completion_tokens=100k
    the total is 200k which would exceed a 200k limit, but compaction
    should NOT trigger because prompt_tokens alone is below threshold.
    """
    usage = TokenUsage(prompt_tokens=100_000, completion_tokens=100_000)
    config = CompactionConfig(auto=True, reserved=20_000)

    result = agent_compaction.should_compact(
        token_usage=usage,
        compaction_config=config,
        model_context_limit=200_000,
    )

    # prompt_tokens=100k < 200k-20k=180k threshold → no compaction
    assert result is False


@pytest.mark.asyncio
async def test_run_compaction_uses_persona_prompt_and_standalone_run(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """run_compaction uses persona prompt and runs summary with empty history."""
    calls: list[dict[str, Any]] = []
    agent: Any = _StubAgent(calls=calls, output="condensed summary")

    system_message = ModelRequest(parts=[UserPromptPart(content="system")])
    message_history = [
        system_message,
        ModelRequest(parts=[UserPromptPart(content="older user message")]),
        ModelResponse(parts=[TextPart(content="assistant reply")]),
        ModelRequest(parts=[UserPromptPart(content="recent user message")]),
    ]

    log_calls: list[tuple[str, str, dict[str, Any]]] = []

    def _fake_append_log(
        instance_id: str,
        event: str,
        details: dict[str, Any],
    ) -> Success[None]:
        log_calls.append((instance_id, event, details))
        return Success(None)

    monkeypatch.setattr(agent_compaction, "append_log", _fake_append_log)

    result = await agent_compaction.run_compaction(
        agent=agent,
        message_history=message_history,
        compaction_config=CompactionConfig(
            auto=True, prune=False, prompt="config prompt"
        ),
        persona_compaction_prompt="persona prompt",
        instance_id="inst_123",
    )

    assert isinstance(result, Success)
    new_history = result.unwrap()

    assert len(calls) == 1
    # Persona prompt is used as the base; the old message history is embedded
    # into the user_prompt so the summary run can be standalone.
    assert calls[0]["user_prompt"].startswith("persona prompt")
    assert "Conversation to summarize" in calls[0]["user_prompt"]
    assert "older user message" in calls[0]["user_prompt"]
    assert "{history}" not in calls[0]["user_prompt"]
    assert calls[0]["kwargs"]["message_history"] == []

    assert new_history[0] is system_message
    assert len(new_history) == 4
    assert isinstance(new_history[1], ModelRequest)
    summary_content = new_history[1].parts[0].content
    assert isinstance(summary_content, str)
    assert "<compaction-summary>" in summary_content
    assert "condensed summary" in summary_content

    assert log_calls == [
        (
            "inst_123",
            "compaction",
            {
                "old_count": 4,
                "new_count": 4,
                "summary_length": len("condensed summary"),
            },
        )
    ]


@pytest.mark.asyncio
async def test_run_compaction_prunes_when_enabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """run_compaction calls prune_tool_outputs when prune=True."""
    calls: list[dict[str, Any]] = []
    agent: Any = _StubAgent(calls=calls, output="summary")

    message_history = [
        ModelRequest(parts=[UserPromptPart(content="system")]),
        ModelRequest(parts=[UserPromptPart(content="older")]),
        ModelResponse(parts=[TextPart(content="recent")]),
    ]

    prune_called = {"value": False}

    def _fake_prune(messages: list[Any], preserve_after_index: int) -> list[Any]:
        prune_called["value"] = True
        assert preserve_after_index >= 1
        return messages

    monkeypatch.setattr(agent_compaction, "prune_tool_outputs", _fake_prune)
    monkeypatch.setattr(
        agent_compaction, "append_log", lambda *_args, **_kwargs: Success(None)
    )

    result = await agent_compaction.run_compaction(
        agent=agent,
        message_history=message_history,
        compaction_config=CompactionConfig(prune=True),
        persona_compaction_prompt=None,
        instance_id="inst_234",
    )

    assert isinstance(result, Success)
    assert prune_called["value"] is True


@pytest.mark.asyncio
async def test_run_compaction_returns_failure_on_llm_error(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """LLM errors are wrapped as InternalError failure."""

    class _FailingAgent:
        async def run(self, *_args: Any, **_kwargs: Any) -> AgentRunResult[str]:
            raise RuntimeError("upstream unavailable")

    message_history = [
        ModelRequest(parts=[UserPromptPart(content="system")]),
        ModelRequest(parts=[UserPromptPart(content="older")]),
    ]

    agent: Any = _FailingAgent()
    # Avoid real filesystem logging on failure path.
    monkeypatch.setattr(agent_compaction, "append_log", lambda *_a, **_k: Success(None))

    result = await agent_compaction.run_compaction(
        agent=agent,
        message_history=message_history,
        compaction_config=CompactionConfig(),
        persona_compaction_prompt=None,
        instance_id="inst_err",
    )

    assert isinstance(result, Failure)
    error = result.failure()
    assert isinstance(error, InternalError)
    assert "Compaction LLM call failed" in error.message
