"""Tool output pruning and budgeting logic for model assembly.

Pure computation functions for pruning old tool call outputs from
PydanticAI message history to reduce context size.

No I/O imports.
"""

import json
from typing import Any

from deal import post, pre
from pydantic_ai.messages import ModelRequest, ModelResponse, ToolReturnPart


PLACEHOLDER = "[Old tool result content cleared]"
HISTORY_PLACEHOLDER_KEY = "placeholder"
TRUNCATED_MARKER = "...[truncated]"


@post(lambda result: isinstance(result, int) and result >= 0)
def _content_size_bytes(value: Any) -> int:
    """Estimate serialized byte size for conservative budget checks."""
    if isinstance(value, str):
        return len(value.encode("utf-8"))
    try:
        return len(
            json.dumps(
                value,
                ensure_ascii=False,
                default=str,
                sort_keys=True,
                separators=(",", ":"),
            ).encode("utf-8")
        )
    except (TypeError, ValueError):
        return len(str(value).encode("utf-8"))


@pre(lambda text, budget_bytes: isinstance(text, str) and budget_bytes >= 0)
@post(lambda result: isinstance(result, str))
def _truncate_text_to_budget(text: str, budget_bytes: int) -> str:
    """Truncate text to a UTF-8 byte budget."""
    if budget_bytes <= 0:
        return ""
    encoded = text.encode("utf-8")
    if len(encoded) <= budget_bytes:
        return text

    marker = TRUNCATED_MARKER
    marker_bytes = marker.encode("utf-8")
    if budget_bytes <= len(marker_bytes):
        return marker_bytes[:budget_bytes].decode("utf-8", errors="ignore")

    keep_bytes = budget_bytes - len(marker_bytes)
    truncated = encoded[:keep_bytes].decode("utf-8", errors="ignore")
    return truncated + marker


@post(lambda result: result is not None)
def _degrade_historical_tool_content(content: Any) -> Any:
    """Degrade older tool content to summary+placeholder when possible."""
    if isinstance(content, dict):
        summary = content.get("summary")
        if summary is not None:
            return {
                "summary": summary,
                HISTORY_PLACEHOLDER_KEY: PLACEHOLDER,
            }
    return PLACEHOLDER


@post(lambda result: isinstance(result, bool))
def _is_historical_placeholder_content(content: Any) -> bool:
    """Return whether a content payload matches historical placeholder policy."""
    if content == PLACEHOLDER:
        return True
    if isinstance(content, dict):
        return content.get(HISTORY_PLACEHOLDER_KEY) == PLACEHOLDER
    return False


@pre(
    lambda content, budget_bytes: (
        isinstance(budget_bytes, int) and budget_bytes >= 0 and content is not Ellipsis
    )
)
@post(lambda result: result is not None)
def _degrade_current_turn_tool_content(content: Any, budget_bytes: int) -> Any:
    """Degrade current-turn tool content under a per-item budget."""
    if budget_bytes <= 0:
        return ""

    if _content_size_bytes(content) <= budget_bytes:
        return content

    if isinstance(content, dict):
        summary = content.get("summary")
        excerpt = content.get("excerpt")
        if summary is not None or excerpt is not None:
            candidate: dict[str, Any] = {}
            if summary is not None:
                candidate["summary"] = summary
            if excerpt is not None:
                candidate["excerpt"] = excerpt
            if _content_size_bytes(candidate) <= budget_bytes:
                return candidate

            if "excerpt" in candidate:
                candidate["excerpt"] = PLACEHOLDER
                if _content_size_bytes(candidate) <= budget_bytes:
                    return candidate

            if "summary" in candidate:
                candidate["summary"] = PLACEHOLDER
                if _content_size_bytes(candidate) <= budget_bytes:
                    return candidate

    text = content if isinstance(content, str) else json.dumps(content, default=str)
    return _truncate_text_to_budget(text, budget_bytes)


@pre(
    lambda part, content: (
        hasattr(part, "tool_name")
        and hasattr(part, "tool_call_id")
        and hasattr(part, "content")
        and content is not Ellipsis
        # Note: content type may differ from original (e.g., dict -> str
        # after truncation fallback). This is intentional — the post-condition
        # validates structural correctness of the result.
    )
)
@post(
    lambda result: (
        hasattr(result, "tool_name")
        and hasattr(result, "tool_call_id")
        and hasattr(result, "content")
    )
)
def _clone_part_with_content(part: Any, content: Any) -> Any:
    """Clone ToolReturnPart while preserving pydantic-ai-required structure."""
    try:
        if hasattr(part, "model_copy"):
            return part.model_copy(update={"content": content})
        return ToolReturnPart(
            tool_name=part.tool_name,
            tool_call_id=part.tool_call_id,
            content=content,
        )
    except (TypeError, ValueError):
        # Keep original part when content cannot be represented in ToolReturnPart.
        return part


@post(
    lambda result: all(
        not isinstance(part, ToolReturnPart) or part.content == PLACEHOLDER
        for part in result
    )
)
def _parts_with_placeholder(parts: list[Any]) -> list[Any]:
    """Replace ToolReturnPart content in a message part list."""
    return [
        _clone_part_with_content(
            part,
            _degrade_historical_tool_content(part.content),
        )
        if isinstance(part, ToolReturnPart)
        else part
        for part in parts
    ]


@pre(
    lambda original_messages, pruned_messages, preserve_after_index: (
        isinstance(original_messages, list)
        and isinstance(pruned_messages, list)
        and isinstance(preserve_after_index, int)
        and preserve_after_index >= 0
    )
)
@post(lambda result: result is None)
def _assert_prune_contract(
    original_messages: list[Any],
    pruned_messages: list[Any],
    preserve_after_index: int,
) -> None:
    """Assert list length and pre-boundary replacement invariants."""
    assert len(pruned_messages) == len(original_messages)

    for idx, original_message in enumerate(original_messages[:preserve_after_index]):
        pruned_message = pruned_messages[idx]
        if not hasattr(original_message, "parts") or not hasattr(
            pruned_message, "parts"
        ):
            continue
        original_parts = original_message.parts
        pruned_parts = pruned_message.parts
        if not isinstance(original_parts, list) or not isinstance(pruned_parts, list):
            continue
        for original_part, pruned_part in zip(original_parts, pruned_parts):
            if isinstance(original_part, ToolReturnPart):
                assert isinstance(pruned_part, ToolReturnPart)
                assert _is_historical_placeholder_content(pruned_part.content)


@pre(
    lambda messages, preserve_after_index, per_call_cap_bytes, per_turn_cap_bytes: (
        isinstance(messages, list)
        and isinstance(preserve_after_index, int)
        and preserve_after_index >= 0
        and preserve_after_index <= len(messages)
        and isinstance(per_call_cap_bytes, int)
        and per_call_cap_bytes >= 1
        and isinstance(per_turn_cap_bytes, int)
        and per_turn_cap_bytes >= 1
    )
)
@post(lambda result: isinstance(result, list))
def assemble_tool_output_history(
    messages: list[Any],
    preserve_after_index: int,
    per_call_cap_bytes: int,
    per_turn_cap_bytes: int,
) -> list[Any]:
    """Assemble message history under tool-output budgets before a model call.

    Strategy:
    - Older messages (index < preserve_after_index): degrade to summary/placeholder.
    - Current-turn messages (index >= preserve_after_index): enforce per-call and
      cumulative per-turn caps, preserving summary+excerpt when possible.

    Doctests:
        >>> from pydantic_ai.messages import ModelRequest, ToolReturnPart
        >>> old = ModelRequest(parts=[ToolReturnPart(content={"summary": "old", "raw": "x" * 200}, tool_name="tree", tool_call_id="1")])
        >>> cur = ModelRequest(parts=[ToolReturnPart(content={"summary": "s", "excerpt": "e", "blob": "y" * 400}, tool_name="tree", tool_call_id="2")])
        >>> assembled = assemble_tool_output_history([old, cur], preserve_after_index=1, per_call_cap_bytes=120, per_turn_cap_bytes=120)
        >>> assembled[0].parts[0].content["summary"]
        'old'
        >>> assembled[0].parts[0].content["placeholder"]
        '[Old tool result content cleared]'
        >>> isinstance(assembled[1].parts[0].content, dict)
        True
        >>> "blob" in assembled[1].parts[0].content
        False
    """
    assembled: list[Any] = []
    turn_budget_used = 0

    for idx, message in enumerate(messages):
        if not hasattr(message, "parts") or not isinstance(message.parts, list):
            assembled.append(message)
            continue

        new_parts: list[Any] = []
        for part in message.parts:
            if not isinstance(part, ToolReturnPart):
                new_parts.append(part)
                continue

            if idx < preserve_after_index:
                degraded = _degrade_historical_tool_content(part.content)
                new_parts.append(_clone_part_with_content(part, degraded))
                continue

            remaining_turn_budget = max(0, per_turn_cap_bytes - turn_budget_used)
            current_budget = min(per_call_cap_bytes, remaining_turn_budget)
            degraded_current = _degrade_current_turn_tool_content(
                part.content,
                current_budget,
            )
            size = _content_size_bytes(degraded_current)
            if size > current_budget:
                degraded_current = _truncate_text_to_budget(
                    str(degraded_current),
                    current_budget,
                )
                size = _content_size_bytes(degraded_current)

            turn_budget_used += size
            new_parts.append(_clone_part_with_content(part, degraded_current))

        if hasattr(message, "model_copy"):
            assembled.append(message.model_copy(update={"parts": new_parts}))
        elif isinstance(message, ModelRequest):
            assembled.append(ModelRequest(parts=new_parts))
        elif isinstance(message, ModelResponse):
            assembled.append(ModelResponse(parts=new_parts))
        else:
            assembled.append(message)

    return assembled


@pre(
    lambda messages, preserve_after_index: (
        isinstance(messages, list)
        and isinstance(preserve_after_index, int)
        and preserve_after_index >= 0
        and preserve_after_index <= len(messages)
    )
)
@post(lambda result: isinstance(result, list))
def prune_tool_outputs(messages: list[Any], preserve_after_index: int) -> list[Any]:
    """Prune old tool call outputs from message history.

    Replaces ToolReturnPart content in the range [0, preserve_after_index)
    with a placeholder string to reduce context size while preserving structure.
    Messages at index >= preserve_after_index are left untouched.

    Args:
        messages: Full message history as list of ModelRequest/ModelResponse.
        preserve_after_index: Index from which to preserve tool outputs.
            All tool results before this index will be replaced.

    Returns:
        New message list with old tool return content replaced.
        Length is guaranteed to match input length (verified in doctests).

    Doctests:
        >>> # Prune tool outputs before index 2
        >>> from pydantic_ai.messages import ModelRequest, ModelResponse, ToolReturnPart, UserPromptPart, TextPart
        >>> messages = [
        ...     ModelRequest(parts=[UserPromptPart(content="user1"), ToolReturnPart(content="old tool result", tool_name="test", tool_call_id="call1")]),
        ...     ModelRequest(parts=[UserPromptPart(content="user2")]),
        ...     ModelRequest(parts=[UserPromptPart(content="user3"), ToolReturnPart(content="new tool result", tool_name="test", tool_call_id="call2")]),
        ... ]
        >>> pruned = prune_tool_outputs(messages, preserve_after_index=2)
        >>> len(pruned) == len(messages)
        True
        >>> # First tool result is replaced
        >>> isinstance(pruned[0].parts[1], ToolReturnPart)
        True
        >>> pruned[0].parts[1].content
        '[Old tool result content cleared]'
        >>> # Second message has no tool result
        >>> len(pruned[1].parts)
        1
        >>> # Third message tool result is preserved
        >>> isinstance(pruned[2].parts[1], ToolReturnPart)
        True
        >>> pruned[2].parts[1].content
        'new tool result'

        >>> # No tool returns to prune
        >>> messages_no_tools = [
        ...     ModelRequest(parts=[UserPromptPart(content="user")]),
        ...     ModelResponse(parts=[TextPart(content="response")]),
        ... ]
        >>> pruned = prune_tool_outputs(messages_no_tools, preserve_after_index=1)
        >>> len(pruned) == 2
        True
        >>> pruned[0].parts[0].content
        'user'

        >>> # Prune all tool outputs (preserve_after_index at end)
        >>> messages_all = [
        ...     ModelRequest(parts=[ToolReturnPart(content="result", tool_name="t", tool_call_id="c")]),
        ... ]
        >>> pruned = prune_tool_outputs(messages_all, preserve_after_index=1)
        >>> pruned[0].parts[0].content
        '[Old tool result content cleared]'

        >>> # Contract: preserve_after_index must be valid
        >>> prune_tool_outputs([], -1)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        PreContractError: ...

        >>> prune_tool_outputs([], 1)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        PreContractError: ...

        >>> prune_tool_outputs([ModelRequest(parts=[])], 2)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        PreContractError: ...
    """
    result_messages: list[Any] = []

    for idx, msg in enumerate(messages):
        # Only prune in the range [0, preserve_after_index)
        if idx < preserve_after_index:
            # Check if this message has parts
            if hasattr(msg, "parts") and isinstance(msg.parts, list):
                # Process parts and replace ToolReturnPart content
                new_parts = _parts_with_placeholder(list(msg.parts))

                # Reconstruct message with new parts
                if isinstance(msg, ModelRequest):
                    result_messages.append(ModelRequest(parts=new_parts))
                elif isinstance(msg, ModelResponse):
                    result_messages.append(ModelResponse(parts=new_parts))
                elif hasattr(msg, "model_copy"):
                    result_messages.append(msg.model_copy(update={"parts": new_parts}))
                else:
                    # Unknown message type is only supported when it has no tool returns.
                    if any(isinstance(part, ToolReturnPart) for part in msg.parts):
                        raise TypeError(
                            "Unsupported message type for tool output pruning: "
                            f"{type(msg).__name__}"
                        )
                    result_messages.append(msg)
            else:
                # No parts or not a standard message, keep as-is
                result_messages.append(msg)
        else:
            # Preserve all messages at index >= preserve_after_index
            result_messages.append(msg)

    _assert_prune_contract(messages, result_messages, preserve_after_index)
    return result_messages


# Re-export for module public API
__all__ = ["assemble_tool_output_history", "prune_tool_outputs"]
