"""Tests for Phase 4 context management: masking, offloading, pressure detection."""

import json
from pathlib import Path

import pytest

from pydantic_ai.messages import (
    ModelRequest,
    ModelResponse,
    UserPromptPart,
    SystemPromptPart,
    TextPart,
    ToolCallPart,
    ToolReturnPart,
)

from autodebug.agent.context import (
    KEEP_VERBATIM_TURNS,
    OFFLOAD_CHAR_THRESHOLD,
    PRESSURE_WARNING_THRESHOLD,
    CHARS_PER_TOKEN,
    _count_user_turns,
    _summarize_tool_result,
    _offload_tool_result,
    _find_tool_call_args,
    mask_old_observations,
    estimate_context_tokens,
    check_context_pressure,
    render_pressure_warning,
    cleanup_offloaded_files,
)


# ---------------------------------------------------------------------------
# Helpers to build synthetic message histories
# ---------------------------------------------------------------------------

def _user_msg(text: str) -> ModelRequest:
    return ModelRequest(parts=[UserPromptPart(content=text)])


def _system_msg(text: str) -> ModelRequest:
    return ModelRequest(parts=[SystemPromptPart(content=text)])


def _tool_return(tool_name: str, content: str, tool_call_id: str = "tc_1") -> ModelRequest:
    return ModelRequest(parts=[ToolReturnPart(tool_name=tool_name, content=content, tool_call_id=tool_call_id)])


def _agent_text(text: str) -> ModelResponse:
    return ModelResponse(parts=[TextPart(content=text)])


def _agent_tool_call(tool_name: str, args: dict, tool_call_id: str = "tc_1") -> ModelResponse:
    return ModelResponse(parts=[ToolCallPart(tool_name=tool_name, args=args, tool_call_id=tool_call_id)])


def _build_multi_turn_history(num_turns: int, content_size: int = 200) -> list:
    """Build a history with N user turns, each with a tool call cycle."""
    history = [_system_msg("You are Hikaflow.")]
    content = "x" * content_size
    for i in range(num_turns):
        tc_id = f"tc_{i}"
        history.append(_user_msg(f"Turn {i}"))
        history.append(_agent_tool_call("read_file", {"path": f"file_{i}.py"}, tool_call_id=tc_id))
        history.append(_tool_return("read_file", content, tool_call_id=tc_id))
        history.append(_agent_text(f"Result for turn {i}"))
    return history


# ---------------------------------------------------------------------------
# _count_user_turns
# ---------------------------------------------------------------------------

class TestCountUserTurns:
    def test_empty_history(self):
        assert _count_user_turns([]) == []

    def test_counts_user_prompts_only(self):
        history = [_system_msg("sys"), _user_msg("hello"), _agent_text("hi")]
        indices = _count_user_turns(history)
        assert len(indices) == 1
        assert indices[0] == 1

    def test_tool_return_only_is_not_a_turn(self):
        history = [
            _user_msg("scan"),
            _agent_tool_call("scan_codebase", {}),
            _tool_return("scan_codebase", "results"),
            _agent_text("done"),
        ]
        indices = _count_user_turns(history)
        assert len(indices) == 1

    def test_multiple_user_turns(self):
        history = _build_multi_turn_history(4)
        indices = _count_user_turns(history)
        assert len(indices) == 4


# ---------------------------------------------------------------------------
# _summarize_tool_result
# ---------------------------------------------------------------------------

class TestSummarizeToolResult:
    def test_read_file_template(self):
        content = "line1\nline2\nline3"
        summary = _summarize_tool_result("read_file", content, args={"path": "api/db.py"})
        assert "[masked]" in summary
        assert "read_file" in summary
        assert "api/db.py" in summary

    def test_default_template(self):
        summary = _summarize_tool_result("unknown_tool", "x" * 500)
        assert "[masked]" in summary
        assert "unknown_tool" in summary
        assert "500" in summary

    def test_offload_path_appended(self):
        summary = _summarize_tool_result("read_file", "content", offload_path=".hikaflow/tmp/abc.txt")
        assert ".hikaflow/tmp/abc.txt" in summary

    def test_string_args_parsed(self):
        summary = _summarize_tool_result("read_file", "x" * 200, args='{"path": "test.py"}')
        assert "test.py" in summary

    def test_bad_json_args_handled(self):
        summary = _summarize_tool_result("read_file", "x" * 200, args="not json")
        assert "[masked]" in summary  # doesn't crash


# ---------------------------------------------------------------------------
# _offload_tool_result
# ---------------------------------------------------------------------------

class TestOffloadToolResult:
    def test_below_threshold_returns_none(self, tmp_path):
        result = _offload_tool_result("short", "read_file", str(tmp_path))
        assert result is None

    def test_above_threshold_creates_file(self, tmp_path):
        big = "x" * (OFFLOAD_CHAR_THRESHOLD + 100)
        result = _offload_tool_result(big, "read_file", str(tmp_path))
        assert result is not None
        assert result.startswith(".hikaflow/tmp/")
        full_path = tmp_path / result
        assert full_path.is_file()
        assert full_path.read_text() == big

    def test_deterministic_filename(self, tmp_path):
        content = "x" * (OFFLOAD_CHAR_THRESHOLD + 100)
        path1 = _offload_tool_result(content, "read_file", str(tmp_path))
        path2 = _offload_tool_result(content, "read_file", str(tmp_path))
        assert path1 == path2


# ---------------------------------------------------------------------------
# mask_old_observations
# ---------------------------------------------------------------------------

class TestMaskOldObservations:
    def test_no_masking_when_few_turns(self, tmp_path):
        history = _build_multi_turn_history(2)
        masked = mask_old_observations(history, str(tmp_path), keep_turns=3)
        assert masked == 0

    def test_masks_old_turns(self, tmp_path):
        history = _build_multi_turn_history(5)
        masked = mask_old_observations(history, str(tmp_path), keep_turns=3)
        assert masked == 2  # turns 0 and 1 masked

    def test_masked_content_starts_with_marker(self, tmp_path):
        history = _build_multi_turn_history(5)
        mask_old_observations(history, str(tmp_path), keep_turns=3)
        for msg in history:
            if isinstance(msg, ModelRequest):
                for part in msg.parts:
                    if isinstance(part, ToolReturnPart) and part.tool_call_id in ("tc_0", "tc_1"):
                        assert part.content.startswith("[masked]")

    def test_recent_turns_not_masked(self, tmp_path):
        history = _build_multi_turn_history(5)
        mask_old_observations(history, str(tmp_path), keep_turns=3)
        for msg in history:
            if isinstance(msg, ModelRequest):
                for part in msg.parts:
                    if isinstance(part, ToolReturnPart) and part.tool_call_id in ("tc_2", "tc_3", "tc_4"):
                        assert not part.content.startswith("[masked]")

    def test_never_masks_user_prompts(self, tmp_path):
        history = _build_multi_turn_history(5)
        mask_old_observations(history, str(tmp_path), keep_turns=3)
        for msg in history:
            if isinstance(msg, ModelRequest):
                for part in msg.parts:
                    if isinstance(part, UserPromptPart):
                        assert not str(part.content).startswith("[masked]")

    def test_never_masks_system_prompts(self, tmp_path):
        history = _build_multi_turn_history(5)
        mask_old_observations(history, str(tmp_path), keep_turns=3)
        for msg in history:
            if isinstance(msg, ModelRequest):
                for part in msg.parts:
                    if isinstance(part, SystemPromptPart):
                        assert not part.content.startswith("[masked]")

    def test_never_masks_agent_text(self, tmp_path):
        history = _build_multi_turn_history(5)
        mask_old_observations(history, str(tmp_path), keep_turns=3)
        for msg in history:
            if isinstance(msg, ModelResponse):
                for part in msg.parts:
                    if isinstance(part, TextPart):
                        assert not part.content.startswith("[masked]")

    def test_idempotent(self, tmp_path):
        history = _build_multi_turn_history(5)
        mask1 = mask_old_observations(history, str(tmp_path), keep_turns=3)
        mask2 = mask_old_observations(history, str(tmp_path), keep_turns=3)
        assert mask1 == 2
        assert mask2 == 0

    def test_offloads_large_results(self, tmp_path):
        # Use unique content per turn so hashes differ
        size = OFFLOAD_CHAR_THRESHOLD + 500
        history = [_system_msg("You are Hikaflow.")]
        for i in range(5):
            tc_id = f"tc_{i}"
            content = f"turn{i}_" + "x" * (size - 6)
            history.append(_user_msg(f"Turn {i}"))
            history.append(_agent_tool_call("read_file", {"path": f"file_{i}.py"}, tool_call_id=tc_id))
            history.append(_tool_return("read_file", content, tool_call_id=tc_id))
            history.append(_agent_text(f"Result for turn {i}"))
        mask_old_observations(history, str(tmp_path), keep_turns=3)
        tmp_dir = tmp_path / ".hikaflow" / "tmp"
        assert tmp_dir.is_dir()
        offloaded = list(tmp_dir.glob("*.txt"))
        assert len(offloaded) == 2  # turns 0 and 1

    def test_skips_tiny_results(self, tmp_path):
        history = _build_multi_turn_history(5, content_size=50)
        masked = mask_old_observations(history, str(tmp_path), keep_turns=3)
        assert masked == 0


# ---------------------------------------------------------------------------
# Context pressure
# ---------------------------------------------------------------------------

class TestContextPressure:
    def test_below_threshold(self):
        history = [_user_msg("hello"), _agent_text("hi")]
        assert check_context_pressure(history) is False

    def test_above_threshold(self):
        big = "x" * (PRESSURE_WARNING_THRESHOLD * CHARS_PER_TOKEN + 1000)
        history = [_user_msg("scan"), _tool_return("scan_codebase", big)]
        assert check_context_pressure(history) is True

    def test_estimate_simple(self):
        history = [_user_msg("hello world")]
        tokens = estimate_context_tokens(history)
        assert tokens == len("hello world") // CHARS_PER_TOKEN

    def test_pressure_warning_content(self):
        warning = render_pressure_warning()
        assert "WARNING" in warning
        assert "memory_insert" in warning or "memory_rethink" in warning


# ---------------------------------------------------------------------------
# Cleanup
# ---------------------------------------------------------------------------

class TestCleanupOffloadedFiles:
    def test_removes_files(self, tmp_path):
        tmp_dir = tmp_path / ".hikaflow" / "tmp"
        tmp_dir.mkdir(parents=True)
        (tmp_dir / "read_file_abc.txt").write_text("content")
        (tmp_dir / "scan_def.txt").write_text("content")
        removed = cleanup_offloaded_files(str(tmp_path))
        assert removed == 2

    def test_handles_missing_dir(self, tmp_path):
        removed = cleanup_offloaded_files(str(tmp_path))
        assert removed == 0


# ---------------------------------------------------------------------------
# _find_tool_call_args
# ---------------------------------------------------------------------------

class TestFindToolCallArgs:
    def test_finds_matching_args(self):
        history = [
            _user_msg("read api/db.py"),
            _agent_tool_call("read_file", {"path": "api/db.py"}, tool_call_id="tc_99"),
            _tool_return("read_file", "contents", tool_call_id="tc_99"),
        ]
        args = _find_tool_call_args(history, "tc_99", max_lookback=3)
        assert args == {"path": "api/db.py"}

    def test_returns_none_when_not_found(self):
        history = [_user_msg("hello")]
        args = _find_tool_call_args(history, "nonexistent", max_lookback=1)
        assert args is None
