"""Policy-wrapping toolset for external MCP tools.

Wraps AbstractToolset to enforce side_effect_policy on tool calls.
For read_only policy, mutating tool calls are blocked with ToolBlockedError
(from core.schemas.errors).

Core zone: @pre/@post contracts + doctests, no I/O imports.
"""

# @invar:allow shell_result: Factory function returns toolset objects, not Result
# @shell_orchestration: Creates wrapped toolsets for policy enforcement

from contextvars import ContextVar, Token
from dataclasses import dataclass
import json
from typing import TYPE_CHECKING, Any

from pydantic_ai.toolsets import AbstractToolset, ToolsetTool

from anima.core.schemas import SideEffectPolicy
from anima.core.tool_policy import (
    ToolClassification,
    ToolDecision,
    classify_tool,
    resolve_tool_policy,
)
from anima.shell.tool_output_telemetry import (
    detect_tool_output_shaping_strategy,
    emit_tool_output_shaped_telemetry,
)

if TYPE_CHECKING:
    pass


MAX_TOOL_RESULT_TEXT_CHARS = 500_000
TRUNCATED_MARKER = "...[truncated]"
PROGRESSIVE_METADATA_KEYS = (
    "next_cursor",
    "continuation_token",
    "has_more",
    "offset",
    "limit",
    "total",
)


@dataclass
class _ToolOutputBudgetState:
    per_call_cap_bytes: int
    per_turn_cap_bytes: int
    used_bytes: int = 0


_TOOL_OUTPUT_BUDGET_STATE: ContextVar[_ToolOutputBudgetState | None] = ContextVar(
    "tool_output_budget_state",
    default=None,
)


# @invar:allow shell_result: Context token scopes turn-local budget state
def begin_tool_output_budget_scope(
    *,
    per_call_cap_bytes: int,
    per_turn_cap_bytes: int,
) -> Token[_ToolOutputBudgetState | None]:
    """Start a turn-local tool-output budget scope.

    This scope is used by PolicyWrappedToolset.call_tool to enforce cumulative
    per-turn budgets for current-turn tool results before model injection.
    """
    return _TOOL_OUTPUT_BUDGET_STATE.set(
        _ToolOutputBudgetState(
            per_call_cap_bytes=max(1, per_call_cap_bytes),
            per_turn_cap_bytes=max(1, per_turn_cap_bytes),
        )
    )


# @invar:allow shell_result: Context token reset returns None by design
def end_tool_output_budget_scope(token: Token[_ToolOutputBudgetState | None]) -> None:
    """End a turn-local tool-output budget scope."""
    _TOOL_OUTPUT_BUDGET_STATE.reset(token)


# @invar:allow shell_result: Scalar byte-size helper for budget accounting
def _content_size_bytes(value: Any) -> int:
    """Estimate serialized payload size in bytes."""
    if isinstance(value, str):
        return len(value.encode("utf-8"))
    try:
        encoded = json.dumps(
            value,
            ensure_ascii=False,
            default=str,
            sort_keys=True,
            separators=(",", ":"),
        ).encode("utf-8")
        return len(encoded)
    except (TypeError, ValueError):
        return len(str(value).encode("utf-8"))


# @invar:allow shell_result: Budget-shaping helper returns truncated text
def _truncate_text_to_budget(text: str, budget_bytes: int) -> str:
    """Truncate text payload to a strict 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 = budget_bytes - len(marker_bytes)
    truncated = encoded[:keep].decode("utf-8", errors="ignore")
    return truncated + marker


# @invar:allow shell_result: Structured shaping helper returns transformed payload
# @invar:allow shell_complexity: Summary/excerpt degradation must branch by remaining budget
def _shape_structured_summary_excerpt(
    content: dict[str, Any], budget_bytes: int
) -> Any:
    """Keep summary+excerpt fields when fitting budget permits."""
    summary = content.get("summary")
    excerpt = content.get("excerpt")
    candidate: dict[str, Any] = {}
    if summary is not None:
        candidate["summary"] = summary
    if excerpt is not None:
        candidate["excerpt"] = excerpt

    if candidate and _content_size_bytes(candidate) <= budget_bytes:
        return candidate

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

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

    return _truncate_text_to_budget(str(content), budget_bytes)


# @invar:allow shell_result: Progressive metadata helper returns shaped dict payload
def _shape_oversized_dict_with_progressive_metadata(
    content: dict[str, Any], budget_bytes: int
) -> Any:
    """Preserve progressive-access metadata when oversized dict payloads are shaped."""
    preserved = {
        key: content[key] for key in PROGRESSIVE_METADATA_KEYS if key in content
    }
    if not preserved:
        return _truncate_text_to_budget(str(content), budget_bytes)

    if _content_size_bytes(preserved) <= budget_bytes:
        return preserved

    trimmed: dict[str, Any] = {}
    for key in PROGRESSIVE_METADATA_KEYS:
        if key not in preserved:
            continue
        value = preserved[key]
        if isinstance(value, str):
            trimmed[key] = _truncate_text_to_budget(value, 32)
        else:
            trimmed[key] = value

    if _content_size_bytes(trimmed) <= budget_bytes:
        return trimmed

    return _truncate_text_to_budget(str(trimmed), budget_bytes)


# @invar:allow shell_result: Current-turn budget helper returns shaped payload object
# @invar:allow shell_complexity: Budget enforcement requires branching on remaining/shape
_HARD_SAFETY_CAP_BYTES = 1_000_000  # 1MB — absolute maximum per tool call


# @invar:allow shell_result: Budget shaping transforms tool output in-place, no Result needed
def _apply_current_turn_budget(raw_result: Any) -> Any:
    """Apply per-call/per-turn budget shaping to current-turn tool output."""
    budget_state = _TOOL_OUTPUT_BUDGET_STATE.get()
    if budget_state is None:
        # No budget scope active (e.g., tool ran in a thread that lost contextvars).
        # Apply a hard safety cap to prevent multi-MB results from reaching the model.
        if isinstance(raw_result, str) and len(raw_result) > _HARD_SAFETY_CAP_BYTES:
            return _truncate_text_to_budget(raw_result, _HARD_SAFETY_CAP_BYTES)
        return raw_result

    remaining_turn = max(0, budget_state.per_turn_cap_bytes - budget_state.used_bytes)
    effective_budget = min(budget_state.per_call_cap_bytes, remaining_turn)
    if effective_budget <= 0:
        return ""

    shaped: Any
    if isinstance(raw_result, dict) and (
        "summary" in raw_result or "excerpt" in raw_result
    ):
        shaped = _shape_structured_summary_excerpt(raw_result, effective_budget)
    elif isinstance(raw_result, dict):
        if _content_size_bytes(raw_result) <= effective_budget:
            shaped = raw_result
        else:
            shaped = _shape_oversized_dict_with_progressive_metadata(
                raw_result, effective_budget
            )
    elif _content_size_bytes(raw_result) <= effective_budget:
        shaped = raw_result
    else:
        text_payload = raw_result if isinstance(raw_result, str) else str(raw_result)
        shaped = _truncate_text_to_budget(text_payload, effective_budget)

    if _content_size_bytes(shaped) > effective_budget:
        shaped = _truncate_text_to_budget(str(shaped), effective_budget)

    budget_state.used_bytes += _content_size_bytes(shaped)
    if budget_state.used_bytes > budget_state.per_turn_cap_bytes:
        budget_state.used_bytes = budget_state.per_turn_cap_bytes
    return shaped


# @invar:allow shell_result: Text-normalization helper returns transformed string
def _truncate_text_content(value: str) -> str:
    """Truncate oversized tool text content with a marker."""
    if len(value) <= MAX_TOOL_RESULT_TEXT_CHARS:
        return value

    suffix = f"\n{TRUNCATED_MARKER}"
    keep_chars = max(0, MAX_TOOL_RESULT_TEXT_CHARS - len(suffix))
    return value[:keep_chars] + suffix


# @invar:allow shell_result: Recursive payload-normalization helper returns transformed object graph
# @invar:allow shell_complexity: Recursive variant handling (str/list/tuple/dict) is required for safe tool payload truncation
def _truncate_tool_result_content(value: Any) -> Any:
    """Recursively truncate oversized string content in tool results."""
    if isinstance(value, str):
        return _truncate_text_content(value)
    if isinstance(value, list):
        return [_truncate_tool_result_content(item) for item in value]
    if isinstance(value, tuple):
        return tuple(_truncate_tool_result_content(item) for item in value)
    if isinstance(value, dict):
        return {key: _truncate_tool_result_content(item) for key, item in value.items()}
    return value


class PolicyWrappedToolset(AbstractToolset[Any]):
    """AbstractToolset wrapper that enforces side_effect_policy on tool calls.

    This wraps an underlying toolset and intercepts call_tool to enforce
    policy decisions. For read_only policy, non-read-only tools are blocked.
    For approval_required policy, non-read-only tools require approval state.

    Attributes:
        _wrapped: The underlying toolset to wrap.
        _policy: The side effect policy to enforce.
        _agent_runner_id: Agent instance ID for approval state lookup.

    Examples:
        >>> # This is tested via integration tests with MockToolset
        >>> pass
    """

    def __init__(
        self,
        wrapped: AbstractToolset[Any],
        policy: SideEffectPolicy,
        agent_runner_id: str,
    ) -> None:
        """Initialize the policy-wrapped toolset.

        Args:
            wrapped: The underlying toolset to wrap.
            policy: The side effect policy to enforce.
            agent_runner_id: Agent instance ID for approval state lookup.
        """
        self._wrapped = wrapped
        self._policy = policy
        self._agent_runner_id = agent_runner_id

    @property
    def id(self) -> str:
        """Return wrapped toolset ID with policy suffix."""
        return f"{self._wrapped.id}:policy-{self._policy.value}"

    async def get_tools(self, ctx: Any) -> dict[str, ToolsetTool[Any]]:
        """Pass through tools from wrapped toolset.

        Args:
            ctx: RunContext from PydanticAI.

        Returns:
            Dict of tool names to ToolsetTool objects.
        """
        return await self._wrapped.get_tools(ctx)

    def _check_approval_state(self, tool_name: str) -> dict[str, Any] | None:
        """Check approval state for approval_required policy.

        TRADE-OFF: Synchronous read_state on each APPROVAL tool call.

        This is a correctness-first design choice. Each tool call reads the
        current state from state.json to check if the agent is in
        'awaiting_approval' state.

        Alternatives considered and rejected:
        1. Cache with invalidation: Complex, risks stale state reads if
           invalidation misses a state transition. Incorrect state could
           allow/block tool calls incorrectly.
        2. Parameterize from AgentRunner: Would require changing
           AbstractToolset.call_tool signature from pydantic-ai, breaking
           abstraction and creating coupling.

        Performance characteristics:
        - state.json is small (< 2KB), very fast to read
        - Only called for APPROVAL decisions (not ALLOW/BLOCK)
        - Frequency bounded by number of tool calls requiring approval
        - AgentRunner already performs state transitions (more expensive writes)

        Escape hatch: If this becomes a performance bottleneck:
        - Add optional state_cache parameter to AgentRunner
        - Pass cached state through tool call context
        - Invalidate cache on state transitions in AgentRunner
        Currently unnecessary: I/O cost is negligible for small file reads.

        Args:
            tool_name: Name of the tool being checked.

        Returns:
            Error dict if approval check fails, None if should proceed.
        """
        from anima.shell.state import clear_approval_response, read_approval_response, read_state
        from returns.result import Failure, Success

        state_result = read_state(self._agent_runner_id)
        if isinstance(state_result, Failure):
            return {
                "error": {
                    "code": "STATE_READ_ERROR",
                    "numeric_code": 500,
                    "message": f"Failed to check approval state: {state_result.failure()}",
                    "details": {"tool_name": tool_name},
                }
            }

        state = state_result.unwrap()

        # Check if agent is in awaiting_approval state
        # This would mean approval was requested but hasn't been resolved yet
        if state.state.value == "awaiting_approval":
            return {
                "error": {
                    "code": "APPROVAL_PENDING",
                    "message": "Approval pending. Wait for approval before retrying this tool.",
                    "details": {"tool_name": tool_name},
                }
            }

        # State is running — check if approval was granted (consume-once)
        if state.state.value == "running":
            response_result = read_approval_response(self._agent_runner_id)
            if isinstance(response_result, Success):
                response = response_result.unwrap()
                if response.get("approved") is True:
                    # Consume the approval so subsequent tool calls require fresh approval
                    clear_approval_response(self._agent_runner_id)
                    return None  # Proceed with tool execution

        # No approval requested yet
        return {
            "error": {
                "code": "APPROVAL_REQUIRED",
                "numeric_code": 331,
                "message": f"Tool {tool_name} requires approval. Use ctl_request_approval to request approval.",
                "details": {"tool_name": tool_name},
            }
        }

    async def call_tool(
        self,
        name: str,
        tool_args: dict[str, Any],
        ctx: Any,
        tool: ToolsetTool[Any],
    ) -> Any:
        """Enforce policy before delegating to wrapped toolset.

        For read_only policy, blocks non-read-only tools.
        For approval_required policy, requires approval state for non-read-only tools.

        Args:
            name: Tool name.
            tool_args: Tool arguments.
            ctx: RunContext from PydanticAI.
            tool: ToolsetTool object.

        Returns:
            Tool result if policy allows execution.

        Raises:
            ToolBlockedError: If policy blocks the tool call.
        """
        from anima.core.schemas.errors import ToolBlockedError

        raw_annotations = getattr(tool, "annotations", None)
        annotations: dict[str, bool] | None = None
        if isinstance(raw_annotations, dict):
            read_only_hint = raw_annotations.get("readOnlyHint")
            destructive_hint = raw_annotations.get("destructiveHint")
            annotations = {
                "readOnlyHint": read_only_hint is True,
                "destructiveHint": destructive_hint is True,
            }

        resolved = resolve_tool_policy(
            self._policy,
            name,
            annotations,
        )
        decision = resolved.decision
        decision_source = resolved.source

        if decision == ToolDecision.BLOCK:
            # Log a distinct warning for unclassifiable tools (default-reject)
            classification = classify_tool(name, annotations)
            if classification == ToolClassification.UNCLASSIFIABLE:
                import logging

                logger = logging.getLogger(__name__)
                logger.warning(
                    f"Unclassifiable tool '{name}' rejected under {self._policy.value} policy. "
                    f"Classification: {classification.value}. "
                    "Default-reject: tools that cannot be classified as read-only are blocked."
                )

            error = ToolBlockedError(name, self._policy.value)
            error_message = error.message
            base_details = error.details if error.details is not None else {}
            error_details = {
                **base_details,
                "source": decision_source,
                "classification": classification.value,
            }
            return {
                "error": {
                    "code": "TOOL_BLOCKED",
                    "numeric_code": 330,
                    "message": error_message,
                    "details": error_details,
                }
            }

        if decision == ToolDecision.APPROVAL:
            # approval_required: check approval state, proceed if approved
            approval_error = self._check_approval_state(name)
            if approval_error is not None:
                return approval_error
            # Approval granted — fall through to execute with budget shaping

        if decision in (ToolDecision.ALLOW, ToolDecision.APPROVAL):
            # Capture budget state BEFORE the tool call. If the tool runs in
            # a thread executor (PydanticAI's run_in_executor for sync tools),
            # contextvars are lost. By capturing here we ensure the budget
            # state is always available for shaping.
            budget_state_before = _TOOL_OUTPUT_BUDGET_STATE.get()

            raw_result = await self._wrapped.call_tool(name, tool_args, ctx, tool)

            # Restore budget state if it was lost during the tool call
            # (e.g., tool ran in a thread that didn't inherit contextvars)
            if budget_state_before is not None and _TOOL_OUTPUT_BUDGET_STATE.get() is None:
                _TOOL_OUTPUT_BUDGET_STATE.set(budget_state_before)

            # Get budget state for telemetry before shaping
            budget_state = _TOOL_OUTPUT_BUDGET_STATE.get()
            per_call_cap = budget_state.per_call_cap_bytes if budget_state else 0

            # Apply budget shaping
            assembled_result = _apply_current_turn_budget(raw_result)

            # Get consumed budget after shaping
            used_after = budget_state.used_bytes if budget_state else 0

            # Determine strategy for telemetry
            strategy = detect_tool_output_shaping_strategy(
                raw_result,
                assembled_result,
                effective_budget=per_call_cap,
                progressive_metadata_keys=PROGRESSIVE_METADATA_KEYS,
                truncated_marker=TRUNCATED_MARKER,
            )

            # Detect source-side pagination capability
            source_side_pagination = isinstance(raw_result, dict) and any(
                key in raw_result for key in PROGRESSIVE_METADATA_KEYS
            )

            # Emit telemetry (sizes only, no content bodies)
            emit_tool_output_shaped_telemetry(
                instance_id=self._agent_runner_id,
                tool_name=name,
                original_content=raw_result,
                shaped_content=assembled_result,
                strategy=strategy,
                per_call_cap_bytes=per_call_cap,
                per_turn_consumed_bytes=used_after,
                source_side_pagination=source_side_pagination,
                content_size_bytes=_content_size_bytes,
            )

            return _truncate_tool_result_content(assembled_result)

        # Should never reach here
        return {
            "error": {
                "code": "UNKNOWN_POLICY",
                "message": f"Unknown policy decision: {decision}",
                "details": {"tool_name": name, "decision": str(decision)},
            }
        }


# @invar:allow shell_result: Factory function returns toolset objects, not Result
# @shell_orchestration: Creates wrapped toolsets for policy enforcement
def wrap_mcp_toolsets_with_policy(
    toolsets: list[AbstractToolset[Any]] | None,
    policy: SideEffectPolicy | None,
    agent_runner_id: str,
) -> list[AbstractToolset[Any]] | None:
    """Wrap MCP toolsets with policy enforcement.

    Args:
        toolsets: List of MCP toolsets to wrap.
        policy: Side effect policy to enforce. None or 'allow' returns unchanged.
        agent_runner_id: Agent instance ID for approval state lookup.

    Returns:
        List of wrapped toolsets, or unchanged when no policy/posture enforcement applies.

    Examples:
        >>> # Integration tests cover the actual wrapping
        >>> pass
    """
    if toolsets is None or not toolsets:
        return toolsets

    if policy is None or policy == SideEffectPolicy.allow:
        # No wrapping needed
        return toolsets

    # Wrap each toolset
    return [
        PolicyWrappedToolset(
            toolset,
            policy if policy is not None else SideEffectPolicy.allow,
            agent_runner_id,
        )
        for toolset in toolsets
    ]


__all__ = [
    "PolicyWrappedToolset",
    "begin_tool_output_budget_scope",
    "end_tool_output_budget_scope",
    "wrap_mcp_toolsets_with_policy",
]
