/headroom/integrations/strands/hooks.py

"""Strands SDK hook provider for Headroom tool output compression.

This module provides HeadroomHookProvider, which implements Strands' HookProvider
interface to intercept tool outputs and compress them using Headroom's SmartCrusher.

Example:
    from strands import Agent
    from headroom.integrations.strands import HeadroomHookProvider

    # Create the hook provider
    hook_provider = HeadroomHookProvider(
        compress_tool_outputs=True,
        min_tokens_to_compress=100,
    )

    # Use with Strands agent
    agent = Agent(hooks=[hook_provider])
    response = agent("Search for documents about AI")

    # Check compression metrics
    print(f"Tokens saved: {hook_provider.total_tokens_saved}")
"""

from __future__ import annotations

import json
import logging
import threading
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any
from uuid import uuid4

# Strands imports - these are optional dependencies
try:
    from strands.hooks import HookProvider, HookRegistry
    from strands.hooks.events import AfterToolCallEvent, BeforeToolCallEvent
    from strands.types.tools import ToolResult

    STRANDS_AVAILABLE = True
except ImportError:
    STRANDS_AVAILABLE = False
    # Type stubs for when strands is not installed
    HookProvider = object  # type: ignore[misc,assignment]
    HookRegistry = object  # type: ignore[misc,assignment]
    AfterToolCallEvent = object  # type: ignore[misc,assignment]
    BeforeToolCallEvent = object  # type: ignore[misc,assignment]
    ToolResult = dict  # type: ignore[misc,assignment]

from headroom import HeadroomConfig
from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig

logger = logging.getLogger(__name__)


def _check_strands_available() -> None:
    """Raise ImportError if Strands is not installed."""
    if not STRANDS_AVAILABLE:
        raise ImportError(
            "Strands SDK is required for this integration. Install with: pip install strands-agents"
        )


def strands_available() -> bool:
    """Check if Strands SDK is installed.

    Returns:
        True if strands-agents package is available.
    """
    return STRANDS_AVAILABLE


@dataclass
class CompressionMetrics:
    """Metrics from a single tool output compression."""

    request_id: str
    timestamp: datetime
    tool_name: str
    tool_use_id: str
    tokens_before: int
    tokens_after: int
    tokens_saved: int
    savings_percent: float
    was_compressed: bool
    skip_reason: str | None = None


@dataclass
class HeadroomHookProvider(HookProvider):  # type: ignore[misc]
    """Strands HookProvider that compresses tool outputs using Headroom.

    This hook provider intercepts tool call results via AfterToolCallEvent
    and applies Headroom's SmartCrusher to compress large outputs, reducing
    token usage while preserving important information.

    The compression is intelligent and preserves:
    - Error items (containing error indicators)
    - Anomalous values (statistical outliers)
    - Items matching the user's query context
    - First/last items for context
    - Structural outliers (rare status values)

    Attributes:
        compress_tool_outputs: Whether to compress tool outputs.
        min_tokens_to_compress: Minimum token count before compression is applied.
        config: Headroom configuration.
        preserve_errors: If True, never compress results with error status.
        total_tokens_saved: Running total of tokens saved across all compressions.
        metrics_history: List of CompressionMetrics from recent compressions.

    Example:
        from strands import Agent
        from headroom.integrations.strands import HeadroomHookProvider

        hook = HeadroomHookProvider(min_tokens_to_compress=50)
        agent = Agent(hooks=[hook])

        # After running agent tasks...
        summary = hook.get_savings_summary()
        print(f"Total saved: {summary['total_tokens_saved']} tokens")
    """

    compress_tool_outputs: bool = True
    min_tokens_to_compress: int = 100
    config: HeadroomConfig | None = field(default=None)
    preserve_errors: bool = True

    # Internal state (not part of dataclass comparison)
    _crusher: SmartCrusher | None = field(default=None, repr=False, compare=False)
    _metrics_history: list[CompressionMetrics] = field(
        default_factory=list, repr=False, compare=False
    )
    _total_tokens_saved: int = field(default=0, repr=False, compare=False)
    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False, compare=False)
    _initialized: bool = field(default=False, repr=False, compare=False)

    def __post_init__(self) -> None:
        """Initialize the hook provider after dataclass construction."""
        _check_strands_available()

        if self.config is None:
            self.config = HeadroomConfig()

        self._initialized = True
        logger.debug(
            "HeadroomHookProvider initialized: compress=%s, min_tokens=%d, preserve_errors=%s",
            self.compress_tool_outputs,
            self.min_tokens_to_compress,
            self.preserve_errors,
        )

    @property
    def crusher(self) -> SmartCrusher:
        """Lazily initialize SmartCrusher (thread-safe).

        Returns:
            The SmartCrusher instance for compression.
        """
        if self._crusher is None:
            with self._lock:
                # Double-check after acquiring lock
                if self._crusher is None:
                    # Use config from HeadroomConfig if available
                    if self.config and self.config.smart_crusher:
                        crusher_config = SmartCrusherConfig(
                            min_tokens_to_crush=self.min_tokens_to_compress,
                            max_items_after_crush=self.config.smart_crusher.max_items_after_crush,
                        )
                    else:
                        crusher_config = SmartCrusherConfig(
                            min_tokens_to_crush=self.min_tokens_to_compress
                        )
                    self._crusher = SmartCrusher(config=crusher_config)
                    logger.debug(
                        "SmartCrusher initialized with min_tokens=%d", self.min_tokens_to_compress
                    )
        return self._crusher

    @property
    def total_tokens_saved(self) -> int:
        """Total tokens saved across all compressions.

        Returns:
            Cumulative token savings.
        """
        return self._total_tokens_saved

    @property
    def metrics_history(self) -> list[CompressionMetrics]:
        """History of compression metrics.

        Returns:
            Copy of the metrics history list.
        """
        return self._metrics_history.copy()

    def register_hooks(self, registry: HookRegistry, **kwargs: Any) -> None:
        """Register hooks with the Strands HookRegistry.

        This method is called by Strands when the hook provider is added
        to an Agent. It registers the compression handler for AfterToolCallEvent.

        Args:
            registry: The Strands HookRegistry to register hooks with.
        """
        if not self.compress_tool_outputs:
            logger.debug("Tool output compression disabled, skipping hook registration")
            return

        # Register the after-tool-call hook for compression
        registry.add_callback(AfterToolCallEvent, self._compress_tool_result)
        logger.info(
            "HeadroomHookProvider registered: compressing tool outputs >= %d tokens",
            self.min_tokens_to_compress,
        )

    def _estimate_tokens(self, text: str) -> int:
        """Estimate token count for text.

        Uses a simple heuristic of ~4 characters per token, which is
        reasonably accurate for English text and JSON content.

        Args:
            text: The text to estimate tokens for.

        Returns:
            Estimated token count.
        """
        if not text:
            return 0
        # ~4 characters per token is a reasonable estimate
        return len(text) // 4

    def _extract_text_content(self, result: ToolResult) -> str:
        """Extract text content from a ToolResult.

        Handles both text and JSON content types in the result.

        Args:
            result: The ToolResult to extract content from.

        Returns:
            String representation of the content.
        """
        content = result.get("content", [])
        if not content:
            return ""

        text_parts = []
        for item in content:
            if isinstance(item, dict):
                if "text" in item:
                    text_parts.append(str(item["text"]))
                elif "json" in item:
                    try:
                        text_parts.append(json.dumps(item["json"], indent=None))
                    except (TypeError, ValueError):
                        text_parts.append(str(item["json"]))
            elif isinstance(item, str):
                text_parts.append(item)

        return "\n".join(text_parts)

    def _update_result_content(self, result: ToolResult, compressed_text: str) -> None:
        """Update the result content with compressed text.

        Modifies the result in place, preserving the original content structure
        (text vs json) where possible.

        Args:
            result: The ToolResult to update (modified in place).
            compressed_text: The compressed content to set.
        """
        content = result.get("content", [])

        if not content:
            # No existing content, create text content
            result["content"] = [{"text": compressed_text}]
            return

        # Try to preserve original structure
        first_item = content[0] if content else None

        if isinstance(first_item, dict):
            if "json" in first_item:
                # Try to parse compressed text back to JSON
                try:
                    parsed = json.loads(compressed_text)
                    result["content"] = [{"json": parsed}]
                except (json.JSONDecodeError, ValueError):
                    # Fall back to text if not valid JSON
                    result["content"] = [{"text": compressed_text}]
            else:
                # Text content
                result["content"] = [{"text": compressed_text}]
        else:
            # Unknown structure, use text
            result["content"] = [{"text": compressed_text}]

    def _compress_tool_result(self, event: AfterToolCallEvent) -> None:
        """Compress tool result content if it exceeds the token threshold.

        This is the main hook handler that intercepts AfterToolCallEvent
        and applies SmartCrusher compression to large tool outputs.

        Args:
            event: The AfterToolCallEvent containing the tool result.
                   The result field is writable and modified in place.
        """
        request_id = str(uuid4())
        result = event.result
        tool_name = event.tool_use.get("name", "unknown")
        tool_use_id = event.tool_use.get("toolUseId", "unknown")

        # Check if compression should be skipped
        skip_reason = self._should_skip_compression(result)
        if skip_reason:
            self._record_metrics(
                request_id=request_id,
                tool_name=tool_name,
                tool_use_id=tool_use_id,
                tokens_before=0,
                tokens_after=0,
                was_compressed=False,
                skip_reason=skip_reason,
            )
            logger.debug(
                "Skipping compression for tool %s (id=%s): %s",
                tool_name,
                tool_use_id,
                skip_reason,
            )
            return

        # Extract content and estimate tokens
        original_text = self._extract_text_content(result)
        tokens_before = self._estimate_tokens(original_text)

        # Check minimum token threshold
        if tokens_before < self.min_tokens_to_compress:
            self._record_metrics(
                request_id=request_id,
                tool_name=tool_name,
                tool_use_id=tool_use_id,
                tokens_before=tokens_before,
                tokens_after=tokens_before,
                was_compressed=False,
                skip_reason=f"below_threshold:{tokens_before}<{self.min_tokens_to_compress}",
            )
            logger.debug(
                "Tool %s output below threshold (%d < %d tokens), skipping compression",
                tool_name,
                tokens_before,
                self.min_tokens_to_compress,
            )
            return

        # Apply compression
        try:
            crush_result = self.crusher.crush(content=original_text, query="")
            compressed_text = crush_result.compressed
            was_modified = crush_result.was_modified
        except Exception as e:
            # Compression failed, keep original
            logger.warning(
                "Compression failed for tool %s (id=%s): %s. Keeping original.",
                tool_name,
                tool_use_id,
                str(e),
            )
            self._record_metrics(
                request_id=request_id,
                tool_name=tool_name,
                tool_use_id=tool_use_id,
                tokens_before=tokens_before,
                tokens_after=tokens_before,
                was_compressed=False,
                skip_reason=f"compression_error:{type(e).__name__}",
            )
            return

        tokens_after = self._estimate_tokens(compressed_text)

        # Only update if compression actually reduced tokens
        if was_modified and tokens_after < tokens_before:
            self._update_result_content(result, compressed_text)
            tokens_saved = tokens_before - tokens_after

            self._record_metrics(
                request_id=request_id,
                tool_name=tool_name,
                tool_use_id=tool_use_id,
                tokens_before=tokens_before,
                tokens_after=tokens_after,
                was_compressed=True,
                skip_reason=None,
            )

            logger.info(
                "Compressed tool %s output: %d -> %d tokens (%.1f%% saved)",
                tool_name,
                tokens_before,
                tokens_after,
                (tokens_saved / tokens_before * 100) if tokens_before > 0 else 0,
            )
        else:
            # Compression didn't help
            self._record_metrics(
                request_id=request_id,
                tool_name=tool_name,
                tool_use_id=tool_use_id,
                tokens_before=tokens_before,
                tokens_after=tokens_before,
                was_compressed=False,
                skip_reason="no_reduction",
            )
            logger.debug(
                "Compression did not reduce tool %s output (%d tokens)",
                tool_name,
                tokens_before,
            )

    def _should_skip_compression(self, result: ToolResult) -> str | None:
        """Check if compression should be skipped for this result.

        Args:
            result: The tool result to check.

        Returns:
            Skip reason string if should skip, None if should compress.
        """
        # Skip if compression is disabled
        if not self.compress_tool_outputs:
            return "compression_disabled"

        # Skip error results if preserve_errors is True
        if self.preserve_errors and result.get("status") == "error":
            return "error_result_preserved"

        # Skip empty results
        content = result.get("content", [])
        if not content:
            return "empty_content"

        return None

    def _record_metrics(
        self,
        request_id: str,
        tool_name: str,
        tool_use_id: str,
        tokens_before: int,
        tokens_after: int,
        was_compressed: bool,
        skip_reason: str | None,
    ) -> None:
        """Record compression metrics (thread-safe).

        Args:
            request_id: Unique ID for this compression request.
            tool_name: Name of the tool that was called.
            tool_use_id: The toolUseId from the result.
            tokens_before: Token count before compression.
            tokens_after: Token count after compression.
            was_compressed: Whether compression was actually applied.
            skip_reason: Reason compression was skipped, if applicable.
        """
        tokens_saved = max(0, tokens_before - tokens_after)
        savings_percent = (tokens_saved / tokens_before * 100) if tokens_before > 0 else 0.0

        metrics = CompressionMetrics(
            request_id=request_id,
            timestamp=datetime.now(timezone.utc),
            tool_name=tool_name,
            tool_use_id=tool_use_id,
            tokens_before=tokens_before,
            tokens_after=tokens_after,
            tokens_saved=tokens_saved,
            savings_percent=savings_percent,
            was_compressed=was_compressed,
            skip_reason=skip_reason,
        )

        with self._lock:
            self._metrics_history.append(metrics)
            if was_compressed:
                self._total_tokens_saved += tokens_saved

            # Keep only last 100 metrics to bound memory
            if len(self._metrics_history) > 100:
                self._metrics_history = self._metrics_history[-100:]

    def get_savings_summary(self) -> dict[str, Any]:
        """Get summary of token savings across all compressions.

        Returns:
            Dictionary with compression statistics including:
            - total_requests: Number of tool outputs processed
            - compressed_requests: Number actually compressed
            - total_tokens_saved: Cumulative tokens saved
            - average_savings_percent: Mean compression ratio
            - total_tokens_before: Sum of all input tokens
            - total_tokens_after: Sum of all output tokens
        """
        if not self._metrics_history:
            return {
                "total_requests": 0,
                "compressed_requests": 0,
                "total_tokens_saved": 0,
                "average_savings_percent": 0.0,
                "total_tokens_before": 0,
                "total_tokens_after": 0,
            }

        compressed_metrics = [m for m in self._metrics_history if m.was_compressed]

        return {
            "total_requests": len(self._metrics_history),
            "compressed_requests": len(compressed_metrics),
            "total_tokens_saved": self._total_tokens_saved,
            "average_savings_percent": (
                sum(m.savings_percent for m in compressed_metrics) / len(compressed_metrics)
                if compressed_metrics
                else 0.0
            ),
            "total_tokens_before": sum(m.tokens_before for m in self._metrics_history),
            "total_tokens_after": sum(m.tokens_after for m in self._metrics_history),
        }

    def reset(self) -> None:
        """Reset all tracked metrics (thread-safe).

        Clears the metrics history and resets the total tokens saved counter.
        Useful for starting fresh measurements or between test runs.
        """
        with self._lock:
            self._metrics_history = []
            self._total_tokens_saved = 0
        logger.debug("HeadroomHookProvider metrics reset")