/headroom/transforms/pipeline.py

"""Transform pipeline orchestration for Headroom SDK."""

from __future__ import annotations

import logging
import os
import threading
import time
from collections.abc import Callable
from contextlib import nullcontext
from typing import TYPE_CHECKING, Any, TypeVar

from ..config import (
    CacheAlignerConfig,
    DiffArtifact,
    HeadroomConfig,
    TransformDiff,
    TransformResult,
    WasteSignals,
)
from ..observability import get_headroom_tracer, get_otel_metrics
from ..tokenizer import Tokenizer
from ..utils import deep_copy_messages
from .base import Transform
from .cache_aligner import CacheAligner
from .content_router import ContentRouter

if TYPE_CHECKING:
    from ..providers.base import Provider

logger = logging.getLogger(__name__)

# Waste-signal detection re-parses the *original* messages for telemetry only
# (it never changes the compression result). On very large transcripts that
# extra parse can take tens of seconds and blow the Anthropic compression
# timeout, making the proxy fail open and discard an already-computed
# compression (#296). Skip the diagnostic above this size to keep the result
# on the critical path.
MAX_WASTE_SIGNAL_DETECTION_TOKENS = 100_000

# A token saving below this is treated as noise — waste-signal detection only
# runs when compression saved more than this many tokens.
_MIN_TOKENS_SAVED_FOR_WASTE_SIGNALS = 100

# OTel GenAI semantic conventions (open-telemetry/semantic-conventions-genai).
# The compression-pipeline span carries this gen_ai.* attribute alongside the
# proprietary headroom.* ones, so Headroom's telemetry groups/filters by the
# standard schema in any OTel-native backend (Grafana, Datadog, etc.). It is a
# string literal rather than an opentelemetry.semconv constant because the gen_ai
# attributes are stability=development (no stable constants are published).
#
# v1 emits only gen_ai.request.model — the one attribute this span can set
# correctly and unconditionally (the model is always known here). The rest are
# deliberately deferred to v2 because this span cannot set them correctly:
#   - gen_ai.operation.name: apply() is shared by many callers (chat, /v1/compress,
#     batch, Gemini countTokens), so no single value is right — it must be threaded
#     from each caller, not hardcoded.
#   - gen_ai.provider.name: Headroom's provider label can't distinguish Bedrock /
#     Gemini from Anthropic / OpenAI at this layer (Bedrock routes via the
#     Anthropic provider).
#   - gen_ai.usage.*: provider-authoritative usage lives on the response path, not
#     this pre-flight compression span (the compressed-input estimate stays under
#     headroom.tokens.after).
GEN_AI_REQUEST_MODEL = "gen_ai.request.model"


_N = TypeVar("_N", int, float)


def _breaker_env(name: str, default: _N, cast: Callable[[str], _N]) -> _N:
    """Parse a circuit-breaker env var, falling back on bad input.

    The breaker is a safety net — a typo'd value must degrade to the
    default with a warning, not crash proxy startup.
    """
    raw = os.environ.get(name)
    if raw is None:
        return default
    try:
        return cast(raw)
    except ValueError:
        logger.warning("Invalid %s=%r; using default %s", name, raw, default)
        return default


class TransformPipeline:
    """
    Orchestrates multiple transforms in the correct order.

    Transform order:
    1. Cache Aligner - normalize prefix for cache hits
    2. Content Router - intelligent content-aware compression (routes to appropriate
       compressor: Kompress for text, SmartCrusher for JSON, CodeCompressor for code, etc.)

    Phase B PR-B1 retired the IntelligentContextManager / RollingWindow
    "drop messages from history" stage. Live-zone-only compression is the
    sole strategy going forward — message-list mutation no longer happens
    in the pipeline.
    """

    def __init__(
        self,
        config: HeadroomConfig | None = None,
        transforms: list[Transform] | None = None,
        provider: Provider | None = None,
    ):
        """
        Initialize pipeline.

        Args:
            config: Headroom configuration.
            transforms: Optional custom transform list (overrides config).
            provider: Provider for model-specific behavior.
        """
        self.config = config or HeadroomConfig()
        self._provider = provider

        if transforms is not None:
            self.transforms = transforms
        else:
            self.transforms = self._build_default_transforms()

        # Circuit breaker (issue #847): after N consecutive pipeline
        # failures, pass messages through untouched for a cooldown window
        # instead of re-running (and re-failing) transforms on every
        # request. Threshold <= 0 disables the breaker.
        self._breaker_threshold = _breaker_env("HEADROOM_PIPELINE_BREAKER_THRESHOLD", 3, int)
        self._breaker_cooldown_s = _breaker_env("HEADROOM_PIPELINE_BREAKER_COOLDOWN_S", 60.0, float)
        self._breaker_lock = threading.Lock()
        self._breaker_failures = 0
        self._breaker_open_until = 0.0

    def _build_default_transforms(self) -> list[Transform]:
        """Build default transform pipeline from config."""
        transforms: list[Transform] = []

        # Order matters!

        # 0. Tool-result interceptors (ast-grep Read outline, etc.) run first
        # so downstream compressors operate on the already-shrunk content.
        # OPT-IN: enable via HeadroomConfig.intercept_tool_results, or for
        # non-config callers (CLI / SDK / tests) the env var
        # HEADROOM_INTERCEPT_ENABLED=1. Off by default while this ships — lets
        # users try it and compare before we make it the default.
        import os as _os

        if getattr(self.config, "intercept_tool_results", False) or _os.environ.get(
            "HEADROOM_INTERCEPT_ENABLED"
        ):
            from headroom.proxy.interceptors import ToolResultInterceptorTransform

            transforms.append(ToolResultInterceptorTransform())

        # 1. Cache Aligner (prefix stabilization)
        if self.config.cache_aligner.enabled:
            transforms.append(CacheAligner(self.config.cache_aligner))

        # 2. Content-aware Compression
        # ContentRouter handles ALL content types intelligently:
        # - JSON arrays -> SmartCrusher
        # - Plain text -> Kompress (ML-based) or passthrough
        # - Code -> CodeCompressor (AST-aware)
        # - Logs -> LogCompressor
        # - Search results -> SearchCompressor
        # - HTML -> HTMLExtractor
        transforms.append(ContentRouter())
        logger.info("Pipeline using ContentRouter for intelligent content-aware compression")

        return transforms

    def _get_tokenizer(self, model: str) -> Tokenizer:
        """Get tokenizer for model.

        Uses provider's tokenizer if available, otherwise falls back to
        the tokenizer registry which auto-detects the best backend per model:
        - OpenAI models: tiktoken (exact)
        - Anthropic models: calibrated estimation (~3.5 chars/token)
        - Open models: HuggingFace tokenizer (if installed)
        - Unknown models: character-based estimation
        """
        if self._provider is not None:
            token_counter = self._provider.get_token_counter(model)
            return Tokenizer(token_counter, model)

        # No provider — use the tokenizer registry (auto-detects per model)
        # TokenCounter from tokenizers and providers have the same interface
        # (count_text, count_messages) but are different Protocol types.
        from headroom.tokenizers import get_tokenizer

        return Tokenizer(get_tokenizer(model), model)  # type: ignore[arg-type]

    def _provider_name(self) -> str | None:
        if self._provider is None:
            return None

        name = getattr(self._provider, "provider_name", None)
        if isinstance(name, str) and name:
            return name

        return self._provider.__class__.__name__.removesuffix("Provider").lower()

    def _breaker_is_open(self) -> bool:
        """True while the circuit breaker cooldown window is active."""
        if self._breaker_threshold <= 0:
            return False
        with self._breaker_lock:
            return time.monotonic() < self._breaker_open_until

    def _breaker_record_failure(self) -> None:
        """Count a pipeline failure; open the breaker at the threshold."""
        if self._breaker_threshold <= 0:
            return
        with self._breaker_lock:
            self._breaker_failures += 1
            if self._breaker_failures >= self._breaker_threshold:
                self._breaker_open_until = time.monotonic() + self._breaker_cooldown_s
                self._breaker_failures = 0
                logger.warning(
                    "Pipeline circuit breaker OPEN after %d consecutive failures; "
                    "passing messages through for %.0fs",
                    self._breaker_threshold,
                    self._breaker_cooldown_s,
                )

    def _breaker_record_success(self) -> None:
        """Reset the consecutive-failure count after a clean run."""
        if self._breaker_threshold <= 0:
            return
        with self._breaker_lock:
            self._breaker_failures = 0

    def apply(
        self,
        messages: list[dict[str, Any]],
        model: str,
        **kwargs: Any,
    ) -> TransformResult:
        """
        Apply all transforms in sequence.

        Args:
            messages: List of messages to transform.
            model: Model name for token counting.
            **kwargs: Additional arguments passed to transforms.
                - model_limit: Context limit override.
                - output_buffer: Output buffer override.
                - tool_profiles: Per-tool compression profiles.
                - request_id: Optional request ID for diff artifact.
                - waste_messages: Optional richer conversion of the same request
                  used for waste-signal detection only (never transformed).

        Returns:
            Combined TransformResult.
        """
        record_metrics = kwargs.pop("record_metrics", True)
        waste_messages = kwargs.pop("waste_messages", None)
        waste_signal_token_limit = int(
            kwargs.pop("waste_signal_token_limit", MAX_WASTE_SIGNAL_DETECTION_TOKENS)
        )
        tokenizer = self._get_tokenizer(model)
        provider_name = self._provider_name()

        # Get model limit from kwargs (should be set by client)
        model_limit = kwargs.get("model_limit")
        if model_limit is None:
            raise ValueError(
                "model_limit is required. Provide it via kwargs or "
                "configure model_context_limits in HeadroomClient."
            )

        # Start with original tokens
        # Circuit breaker open — pass through untouched (issue #847).
        if self._breaker_is_open():
            passthrough_tokens = tokenizer.count_messages(messages)
            return TransformResult(
                messages=messages,
                tokens_before=passthrough_tokens,
                tokens_after=passthrough_tokens,
                transforms_applied=["pipeline:circuit_open"],
            )

        t_count = time.perf_counter()
        tokens_before = tokenizer.count_messages(messages)
        count_ms = (time.perf_counter() - t_count) * 1000

        logger.debug(
            "Pipeline starting: %d messages, %d tokens, model=%s",
            len(messages),
            tokens_before,
            model,
        )

        tracer = get_headroom_tracer()
        span_attributes: dict[str, Any] = {
            "headroom.model": model,
            "headroom.provider": provider_name or "unknown",
            "headroom.message_count": len(messages),
            "headroom.tokens.before": tokens_before,
        }
        # OTel GenAI semconv request descriptor — emitted alongside headroom.* so
        # the span is groupable by the standard schema (v1: model only).
        if model:
            span_attributes[GEN_AI_REQUEST_MODEL] = model
        pipeline_span_context = (
            tracer.start_as_current_span(
                "headroom.compression.pipeline",
                attributes=span_attributes,
            )
            if record_metrics
            else nullcontext()
        )

        with pipeline_span_context as pipeline_span:
            # Track all transforms applied
            all_transforms: list[str] = []
            all_markers: list[str] = []
            all_warnings: list[str] = []
            all_timing: dict[str, float] = {}  # transform_name → ms

            # Track transform diffs if enabled
            transform_diffs: list[TransformDiff] = []
            generate_diff = self.config.generate_diff_artifact

            t_copy = time.perf_counter()
            current_messages = deep_copy_messages(messages)
            copy_ms = (time.perf_counter() - t_copy) * 1000

            all_timing["_deep_copy"] = copy_ms
            all_timing["_initial_token_count"] = count_ms

            pipeline_start = time.perf_counter()

            request_id = kwargs.get("request_id", "")
            log_prefix = f"[{request_id}] " if request_id else ""

            frozen_count = kwargs.get("frozen_message_count", 0)
            if frozen_count > 0:
                logger.info(
                    "%sPipeline: freezing first %d/%d messages (prefix cached by provider)",
                    log_prefix,
                    frozen_count,
                    len(messages),
                )

            for transform in self.transforms:
                # Check if transform should run
                if not transform.should_apply(current_messages, tokenizer, **kwargs):
                    continue

                transform_span_context = (
                    tracer.start_as_current_span(
                        "headroom.compression.transform",
                        attributes={
                            "headroom.model": model,
                            "headroom.provider": provider_name or "unknown",
                            "headroom.transform": transform.name,
                        },
                    )
                    if record_metrics
                    else nullcontext()
                )

                with transform_span_context as transform_span:
                    # Time the transform
                    t0 = time.perf_counter()
                    try:
                        result = transform.apply(current_messages, tokenizer, **kwargs)
                    except Exception:
                        self._breaker_record_failure()
                        raise
                    duration_ms = (time.perf_counter() - t0) * 1000

                    # Update messages for next transform
                    current_messages = result.messages

                    # Use token counts reported by the transform itself — avoids
                    # redundant O(N) recount of the full message list after each step.
                    tokens_before_transform = result.tokens_before
                    tokens_after_transform = result.tokens_after

                    if transform_span is not None and transform_span.is_recording():
                        transform_span.set_attribute(
                            "headroom.tokens.before", tokens_before_transform
                        )
                        transform_span.set_attribute(
                            "headroom.tokens.after", tokens_after_transform
                        )
                        transform_span.set_attribute(
                            "headroom.tokens.saved",
                            tokens_before_transform - tokens_after_transform,
                        )
                        transform_span.set_attribute("headroom.duration_ms", duration_ms)
                        transform_span.set_attribute(
                            "headroom.transforms_applied",
                            len(result.transforms_applied),
                        )

                    # Accumulate results
                    all_transforms.extend(result.transforms_applied)
                    all_markers.extend(result.markers_inserted)
                    all_warnings.extend(result.warnings)
                    all_timing[transform.name] = duration_ms

                    # Merge sub-transform timing (e.g. ContentRouter's per-compressor breakdown)
                    if result.timing:
                        all_timing.update(result.timing)

                    # Log transform results
                    if result.transforms_applied:
                        logger.info(
                            "Transform %s: %d -> %d tokens (saved %d) [%.1fms]",
                            transform.name,
                            tokens_before_transform,
                            tokens_after_transform,
                            tokens_before_transform - tokens_after_transform,
                            duration_ms,
                        )
                    else:
                        logger.debug(
                            "Transform %s: no changes [%.1fms]", transform.name, duration_ms
                        )

                    # Record diff if enabled
                    if generate_diff:
                        transform_diffs.append(
                            TransformDiff(
                                transform_name=transform.name,
                                tokens_before=tokens_before_transform,
                                tokens_after=tokens_after_transform,
                                tokens_saved=tokens_before_transform - tokens_after_transform,
                                details=", ".join(result.transforms_applied)
                                if result.transforms_applied
                                else "",
                                duration_ms=duration_ms,
                            )
                        )

            # All transforms ran without raising — reset the breaker.
            self._breaker_record_success()

            # Single final token count — the only full recount in the pipeline.
            # Earlier per-transform counts come from each transform's own result.
            t_final_count = time.perf_counter()
            tokens_after = tokenizer.count_messages(current_messages)
            all_timing["_final_token_count"] = (time.perf_counter() - t_final_count) * 1000

            pipeline_ms = (time.perf_counter() - pipeline_start) * 1000
            all_timing["pipeline_total"] = pipeline_ms

            # Log pipeline summary
            total_saved = tokens_before - tokens_after
            timing_parts = " ".join(f"{k}={v:.0f}ms" for k, v in all_timing.items())
            if total_saved > 0:
                logger.info(
                    "%sPipeline complete: %d -> %d tokens (saved %d, %.1f%% reduction) [%s]",
                    log_prefix,
                    tokens_before,
                    tokens_after,
                    total_saved,
                    (total_saved / tokens_before * 100) if tokens_before > 0 else 0,
                    timing_parts,
                )
            else:
                logger.debug("%sPipeline complete: no token savings [%s]", log_prefix, timing_parts)

            # Build diff artifact if enabled
            diff_artifact = None
            if generate_diff:
                diff_artifact = DiffArtifact(
                    request_id=kwargs.get("request_id", ""),
                    original_tokens=tokens_before,
                    optimized_tokens=tokens_after,
                    total_tokens_saved=tokens_before - tokens_after,
                    transforms=transform_diffs,
                )

            # Detect waste signals in original messages (only when significant
            # compression). Handlers whose wire format carries tool output the
            # message conversion drops (e.g. Gemini functionResponse parts, #819)
            # pass a richer waste_messages list that is parsed instead — it is
            # telemetry-only and never transformed.
            waste_signals: WasteSignals | None = None
            saved_enough = (
                tokens_before > tokens_after
                and (tokens_before - tokens_after) > _MIN_TOKENS_SAVED_FOR_WASTE_SIGNALS
            )
            if saved_enough and tokens_before > waste_signal_token_limit:
                # Telemetry-only re-parse would risk the compression timeout on a
                # request this large (#296); skip it and keep the result.
                logger.debug(
                    "%sSkipping waste-signal detection for %d-token request "
                    "(limit=%d) to keep the compression result on the critical path",
                    log_prefix,
                    tokens_before,
                    waste_signal_token_limit,
                )
            elif saved_enough:
                try:
                    from ..parser import parse_messages

                    # current_messages (the post-transform copy) enables reread
                    # attribution: repeats whose first serve was markerized by
                    # this pipeline run count into reread_compressed_tokens
                    # (#899). The length guard in parse_messages makes the
                    # waste_messages path (different indexing) a safe no-op.
                    _, _, waste_signals = parse_messages(
                        waste_messages or messages,
                        tokenizer,
                        compressed_messages=current_messages,
                    )
                    if waste_signals.total() == 0:
                        waste_signals = None
                except Exception:
                    pass

            if pipeline_span is not None and pipeline_span.is_recording():
                pipeline_span.set_attribute("headroom.tokens.after", tokens_after)
                pipeline_span.set_attribute("headroom.tokens.saved", total_saved)
                pipeline_span.set_attribute("headroom.duration_ms", pipeline_ms)
                pipeline_span.set_attribute("headroom.transforms_applied", len(all_transforms))
                pipeline_span.set_attribute("headroom.warnings", len(all_warnings))

            if record_metrics:
                get_otel_metrics().record_pipeline_run(
                    model=model,
                    provider=provider_name,
                    tokens_before=tokens_before,
                    tokens_after=tokens_after,
                    duration_ms=pipeline_ms,
                    timing=all_timing,
                    transforms_applied=all_transforms,
                    waste_signals=waste_signals.to_dict() if waste_signals is not None else None,
                )

        return TransformResult(
            messages=current_messages,
            tokens_before=tokens_before,
            tokens_after=tokens_after,
            transforms_applied=all_transforms,
            markers_inserted=all_markers,
            warnings=all_warnings,
            diff_artifact=diff_artifact,
            timing=all_timing,
            waste_signals=waste_signals,
        )

    def simulate(
        self,
        messages: list[dict[str, Any]],
        model: str,
        **kwargs: Any,
    ) -> TransformResult:
        """
        Simulate transforms without modifying messages.

        Same as apply() but returns what WOULD happen.

        Args:
            messages: List of messages.
            model: Model name.
            **kwargs: Additional arguments.

        Returns:
            TransformResult with simulated changes.
        """
        # apply() already works on a copy, so this is safe
        return self.apply(messages, model, record_metrics=False, **kwargs)


def create_pipeline(
    cache_aligner_config: CacheAlignerConfig | None = None,
) -> TransformPipeline:
    """
    Create a pipeline with specific configurations.

    Args:
        cache_aligner_config: Cache aligner configuration.

    Returns:
        Configured TransformPipeline.
    """
    config = HeadroomConfig()

    if cache_aligner_config is not None:
        config.cache_aligner = cache_aligner_config

    return TransformPipeline(config)