nlproxy/service/compression.py

"""
Orchestration service for the complete prompt compression pipeline.

This module coordinates all compression stages:
1. Prompt shielding (entity protection, restriction extraction)
2. Semantic segmentation and embedding generation
3. Clustering-based sentence compression
4. Reconstruction with token optimization
5. Safety validation and intent preservation

Mathematical Foundations
------------------------
1. Pipeline Composition:
    Given input prompt P and configuration C:
        Output = Safety(Reconstruct(Compress(Segment(Shield(P, C)))))
    Each stage is a function fᵢ: Xᵢ → Xᵢ₊₁ with well-defined contracts.

2. Batch Processing Parallelism:
    For batch of N prompts with M workers:
        Time ≈ O(max(⌈N/M⌉ · T_stage)) where T_stage = slowest stage latency
    Amdahl's Law: Speedup ≤ 1 / (S + P/N) where S=serial fraction, P=parallel

3. Embedding Cache Efficiency:
    Hit rate H = |cached| / |total| ∈ [0, 1]
    Expected latency: E[T] = H·T_cache + (1-H)·T_compute
    Where T_cache ≈ O(1) lookup, T_compute = embedding inference time

4. Aggressiveness Adaptation:
    effective_agg = base_agg × (1 + α·nli_active) where α = 0.3
    Higher NLI confidence allows more aggressive compression safely.

References
----------
[1] Amdahl, G. M. (1967). Validity of the single processor approach to 
    achieving large scale computing capabilities. AFIPS Conference.

[2] Reimers, N., & Gurevych, I. (2019). Sentence-BERT: Sentence embeddings 
    using Siamese BERT-networks. EMNLP-IJCNLP 2019.
    https://github.com/UKPLab/sentence-transformers

[3] Sennrich, R., Haddow, B., & Birch, A. (2016). Neural Machine Translation 
    of Rare Words with Subword Units. ACL 2016.
    https://github.com/openai/tiktoken

Performance Characteristics
---------------------------
- compress_batch(): O(N · T_pipeline / M) with M workers, N prompts
- Typical per-prompt latency: 100-500ms (CPU), 50-200ms (GPU)
- Memory: O(B · d) for batch embeddings, B=batch size, d=embedding_dim
- Cache hit rate: 60-90% typical for repetitive prompts

Author: IntelliDeep Labs Team
License: BSL 1.1
"""

from __future__ import annotations

import asyncio
import hashlib
import json
import logging
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any, Dict, List, Optional

import numpy as np

# Import core components from sibling modules

from nlproxy.core.restriction import Restriction
from nlproxy.core.shield import PromptShield, ShieldResult, DomainMode
from nlproxy.core.segmenter import SemanticSegmenter
from nlproxy.core.compressor import SemanticCompressor
from nlproxy.core.reconstructor import PromptReconstructor, ReconstructionResult
from nlproxy.core.safety import SafetyChecker, SafetyReport
from nlproxy.cache.semantic_cache import SemanticLLMCache
from nlproxy.utils.constants import AGGRESSIVENESS_MAP

# Optional Redis import for distributed caching
try:
    from redis import Redis
    _REDIS_AVAILABLE = True
except ImportError:
    _REDIS_AVAILABLE = False
    Redis = None  # type: ignore

logger = logging.getLogger(__name__)


class CompressionService:
    """
    Orchestrates the complete prompt compression pipeline.

    This service coordinates all stages of prompt processing:
    1. Shielding: Protect entities, extract restrictions
    2. Segmentation: Split text into sentences, generate embeddings
    3. Compression: Cluster similar sentences, select representatives
    4. Reconstruction: Re-inject entities, optimize tokens, compute metrics
    5. Safety: Validate intent preservation, enforce constraints

    Key Features
    ------------
    - Batch processing with thread pool parallelism
    - Multi-level caching: shield results, embeddings, semantic responses
    - Adaptive aggressiveness based on domain mode and NLI confidence
    - Async support for non-blocking operation in high-concurrency servers
    - Configurable privacy mode for PII handling

    Pipeline Architecture
    ---------------------
    Input: List[str] prompts + configuration
            ↓
    [Shield] → Protect entities, extract restrictions
            ↓
    [Segment] → Split sentences, compute embeddings (cached)
            ↓
    [Compress] → Cluster sentences, select centroids
            ↓
    [Reconstruct] → Re-inject entities, compute token metrics
            ↓
    [Safety] → Validate constraints, re-insert if needed
            ↓
    Output: List[Dict] with compressed text and metrics

    Usage Example
    -------------
    >>> service = CompressionService(
    ...     use_cache=True,
    ...     redis_url="redis://localhost:6379",
    ...     privacy_mode=True
    ... )
    >>> results = service.compress_batch(
    ...     texts=["Hello world", "Another prompt"],
    ...     mode="code",
    ...     aggressiveness=0.3
    ... )
    >>> for r in results:
    ...     print(f"Saved {r['tokens_saved']} tokens")
    """

    # Aggressiveness presets are centralized in shared constants for consistent tuning.

    # Configuration defaults
    _DEFAULT_MODEL_NAME: str = "all-MiniLM-L6-v2"
    _DEFAULT_EMBEDDING_DIM: int = 384
    _DEFAULT_BATCH_SIZE: int = 128
    _DEFAULT_MAX_SEQ_LENGTH: int = 256
    _DEFAULT_BASE_AGGRESSIVENESS: float = 0.2
    _DEFAULT_NLI_ADAPTATION_FACTOR: float = 0.3
    _DEFAULT_THREAD_POOL_WORKERS: int = 8


    def __init__(
        self,
        use_cache: bool = True,
        device: Optional[str] = None,
        redis_url: Optional[str] = None,
        nli_refinement_fn: Optional = None,
        privacy_mode: bool = False,
        models_dir: Optional[Path] = None,
        llm_default_model: Optional[str] = None,
        thread_pool_workers: Optional[int] = None,
    ) -> None:
        """
        Initialize the CompressionService with all pipeline components.

        Parameters
        ----------
        use_cache : bool, optional
            Enable in-memory caching for shield results and embeddings.
        device : Optional[str], optional
            Device for embedding model ("cuda", "cpu", or None for auto-detect).
        redis_url : Optional[str], optional
            Redis connection URL for distributed semantic caching.
            If None, semantic cache is disabled.
        nli_refinement_fn : Optional, optional
            NLI inference function for restriction refinement.
            Typically obtained from PostLLMVerifier.get_nli_check_function().
        privacy_mode : bool, optional
            Enable strict PII handling: never expose protected entities.
        models_dir : Optional[Path], optional
            Directory containing pre-downloaded models (default: "models").
            Required for embedding and NLI models.

        Raises
        ------
        ImportError
            If redis_url is provided but redis-py is not installed.
        FileNotFoundError
            If required models are not found in models_dir.

        Complexity
        ----------
        Time: O(1) initialization + O(T_load) for model loading
        Space: O(C) for cache storage if enabled, C = cache capacity
        """
        # Validate Redis availability if requested
        if redis_url and not _REDIS_AVAILABLE:
            raise ImportError(
                "Redis URL provided but redis-py not installed. "
                "Install with: pip install redis"
            )

        # Resolve models directory
        self.models_dir = models_dir or Path("nlproxy") / "models"

        # Thread pool sizing can be overridden via environment variable
        if thread_pool_workers is not None:
            self.thread_pool_workers = thread_pool_workers
        else:
            env_workers = os.getenv("NLPROXY_COMPRESSION_WORKERS", "")
            try:
                self.thread_pool_workers = int(env_workers)
            except ValueError:
                self.thread_pool_workers = self._DEFAULT_THREAD_POOL_WORKERS

        if self.thread_pool_workers <= 0:
            self.thread_pool_workers = self._DEFAULT_THREAD_POOL_WORKERS
        self.executor = ThreadPoolExecutor(max_workers=self.thread_pool_workers)
        self.privacy_mode = privacy_mode

        # Initialize pipeline components
        self.shield = PromptShield(mode=DomainMode.GENERAL)

        self.segmenter = SemanticSegmenter(
            model_name=self._DEFAULT_MODEL_NAME,
            device=device,
            batch_size=self._DEFAULT_BATCH_SIZE,
            max_seq_length=self._DEFAULT_MAX_SEQ_LENGTH,
            models_dir=self.models_dir,
        )

        self.compressor = SemanticCompressor(
            aggressiveness=self._DEFAULT_BASE_AGGRESSIVENESS
        )

        # Use provided default LLM model for token counting/pricing when available
        model_for_reconstructor = llm_default_model or "gpt-4"
        self.reconstructor = PromptReconstructor(model_name=model_for_reconstructor)

        self.safety = SafetyChecker(
            mode="general",
            models_dir=self.models_dir,
        )

        # Caching configuration
        self.use_emb_cache = use_cache
        self.emb_cache: Optional[Dict[str, bytes]] = {} if use_cache else None
        self.shield_cache: Optional[Dict[str, str]] = {} if use_cache else None

        # Redis client for distributed caching (optional)
        self.redis: Optional[Redis] = None
        if redis_url and _REDIS_AVAILABLE:
            self.redis = Redis.from_url(redis_url, decode_responses=True)
            try:
                self.redis.ping()  # Test connection
                logger.info(f"Connected to Redis at {redis_url}")
            except Exception as e:
                logger.warning(f"Redis connection failed: {e}")
                self.redis = None

        # Semantic cache for response deduplication (optional)
        if redis_url and self.redis:
            self.semantic_cache = SemanticLLMCache(
                redis_url=redis_url,
                dimension=self._DEFAULT_EMBEDDING_DIM,
                #models_dir=self.models_dir,
            )
            logger.info("Semantic cache initialized")
        else:
            self.semantic_cache = None
            logger.debug("Semantic cache disabled (Redis not configured)")

        # NLI refinement function for semantic restriction validation
        self.nli_refinement_fn = nli_refinement_fn

        # Privacy mode: controls entity re-injection behavior
        self.privacy_mode = privacy_mode

        # Optional post-processing components (set externally)
        self.post_verifier:Optional = None
        self.response_corrector:Optional = None

        logger.info(
            f"CompressionService initialized: cache={use_cache}, "
            f"redis={redis_url is not None}, privacy={privacy_mode}"
        )


    def _hash_text(self, text: str) -> str:
        """
        Compute SHA-256 hash of text for cache key generation.

        Parameters
        ----------
        text : str
            Input text to hash.

        Returns
        -------
        str
            Hexadecimal SHA-256 hash (64 characters).

        Note
        ----
        SHA-256 provides collision probability < 2⁻¹²⁸ for practical inputs,
        sufficient for cache key uniqueness.
        """
        return hashlib.sha256(text.encode("utf-8")).hexdigest()


    def _cache_shield(
        self,
        text: str,
        shield_result: ShieldResult,
        mode: str = "general",
        privacy_mode: bool = False,
    ) -> None:
        """
        Store shield result in cache (Redis or in-memory).

        Parameters
        ----------
        text : str
            Original prompt text (used as cache key).
        shield_result : ShieldResult
            Result from PromptShield to cache.
        mode : str, optional
            Domain mode for pattern selection.
        privacy_mode : bool, optional
            Whether privacy mode is active.

        Complexity
        ----------
        Time: O(1) for in-memory, O(L) for Redis where L = serialized size
        Space: O(L) for cached entry
        """
        key = f"shield:{mode}:{privacy_mode}:{self._hash_text(text)}"
        data = json.dumps(shield_result.to_cache_dict())

        if self.redis:
            # Redis: set with optional TTL (default 1 hour)
            self.redis.setex(key, 3600, data)
        elif self.shield_cache is not None:
            # In-memory fallback
            self.shield_cache[key] = data


    def _get_cached_shield(
        self,
        text: str,
        mode: str = "general",
        privacy_mode: bool = False,
    ) -> Optional[ShieldResult]:
        """
        Retrieve shield result from cache if available.

        Parameters
        ----------
        text : str
            Original prompt text to look up.
        mode : str, optional
            Domain mode for pattern selection.
        privacy_mode : bool, optional
            Whether privacy mode is active.

        Returns
        -------
        Optional[ShieldResult]
            Cached result if found, None otherwise.

        Complexity
        ----------
        Time: O(1) for in-memory, O(L) for Redis deserialization
        Space: O(1) additional
        """
        key = f"shield:{mode}:{privacy_mode}:{self._hash_text(text)}"
        data: Optional[str] = None

        if self.redis:
            data = self.redis.get(key)
        elif self.shield_cache is not None:
            data = self.shield_cache.get(key)

        if data:
            # Handle bytes from Redis
            if isinstance(data, bytes):
                data = data.decode("utf-8")
            return ShieldResult.from_cache_dict(json.loads(data))

        return None


    def _shield_with_cache(
        self,
        text: str,
        manual_restrictions: Optional[List[Restriction]] = None,
        mode: str = "general",
        privacy_mode: Optional[bool] = None,
    ) -> ShieldResult:
        """
        Shield prompt with cache lookup fallback.

        Parameters
        ----------
        text : str
            Prompt to shield.
        manual_restrictions : Optional[List[Restriction]], optional
            Additional restrictions to enforce.
        mode : str, optional
            Domain mode for pattern selection.
        privacy_mode : Optional[bool], optional
            Override for PII handling.

        Returns
        -------
        ShieldResult
            Shielding result (from cache or freshly computed).

        Complexity
        ----------
        Time: O(1) cache hit, O(T_shield) cache miss
        where T_shield = shielding pipeline latency (~10-50ms)
        """
        effective_privacy = (
            privacy_mode if privacy_mode is not None else self.privacy_mode
        )
        # Try cache first
        cached = self._get_cached_shield(text, mode, effective_privacy)
        if cached:
            logger.debug(f"Shield cache hit for text hash {self._hash_text(text)[:16]}...")
            return cached

        # Compute fresh result
        result = self.shield.shield(
            text,
            manual_restrictions=manual_restrictions,
            nli_refinement_fn=self.nli_refinement_fn,
            privacy_mode=effective_privacy,
            mode_override=mode,
        )

        # Store in cache for future requests
        self._cache_shield(text, result, mode, effective_privacy)
        return result


    def _compute_effective_aggressiveness(
        self,
        aggressiveness: Optional[float],
        mode: str,
        nli_active: bool,
    ) -> float:
        """
        Compute effective aggressiveness considering mode and NLI status.

        Formula:
            effective = (explicit or mode_preset) × (1 + α·nli_active)
            where α = 0.3 (adaptation factor)

        Parameters
        ----------
        aggressiveness : Optional[float]
            Explicit aggressiveness value, or None for mode default.
        mode : str
            Domain mode for preset lookup.
        nli_active : bool
            Whether NLI-based validation is enabled.

        Returns
        -------
        float
            Effective aggressiveness ∈ [0, 1].
        """
        # Resolve base aggressiveness: explicit > mode preset > default
        if aggressiveness is not None:
            base_agg = aggressiveness
        else:
            base_agg = AGGRESSIVENESS_MAP.get(mode, self._DEFAULT_BASE_AGGRESSIVENESS)

        # Adapt for NLI: higher confidence allows more aggressive compression
        if nli_active:
            effective_agg = min(1.0, base_agg * (1 + self._DEFAULT_NLI_ADAPTATION_FACTOR))
            logger.debug(f"NLI adaptation: {base_agg:.2f} → {effective_agg:.2f}")
        else:
            effective_agg = base_agg

        return effective_agg


    def _collect_all_sentences(
        self,
        shield_results: List[ShieldResult],
        language: Optional[str],
    ) -> tuple[List[str], List[tuple[int, List[str]]]]:
        """
        Extract all sentences from shielded texts for batch embedding.

        Parameters
        ----------
        shield_results : List[ShieldResult]
            Results from shielding stage.
        language : Optional[str]
            Language code for segmentation (auto-detect if None).

        Returns
        -------
        Tuple[List[str], List[Tuple[int, List[str]]]]
            - Flat list of all sentences for embedding
            - Mapping: (original_index, sentences_for_that_prompt)
        """
        all_sents: List[str] = []
        sent_map: List[tuple[int, List[str]]] = []

        for i, shielded in enumerate(shield_results):
            sents = self.segmenter.split_sentences(
                shielded.shielded_text, language=language
            )
            sent_map.append((i, sents))
            all_sents.extend(sents)

        return all_sents, sent_map


    def _encode_with_cache(
        self,
        sentences: List[str],
    ) -> Dict[int, np.ndarray]:
        """
        Generate embeddings with in-memory caching.

        Parameters
        ----------
        sentences : List[str]
            Sentences to encode.

        Returns
        -------
        Dict[int, np.ndarray]
            Mapping: sentence_index → embedding vector.

        Complexity
        ----------
        Time: O(U · T_encode + C · T_lookup)
            where U = uncached count, C = cached count
        Space: O(U · d) for new embeddings, d = embedding_dim
        """
        sent_emb: Dict[int, np.ndarray] = {}

        if not self.use_emb_cache or self.emb_cache is None:
            # No caching: encode all at once
            all_embs = self.segmenter.encode_batch(sentences)
            return {i: emb for i, emb in enumerate(all_embs)}

        # Separate cached vs uncached sentences
        uncached: List[str] = []
        uncached_idx: List[int] = []

        for idx, sent in enumerate(sentences):
            key = self._hash_text(sent)
            cached_bytes = self.emb_cache.get(key)

            if cached_bytes is not None:
                # Load from cache
                sent_emb[idx] = np.frombuffer(cached_bytes, dtype=np.float32)
            else:
                uncached.append(sent)
                uncached_idx.append(idx)

        # Encode uncached sentences
        if uncached:
            new_embs = self.segmenter.encode_batch(uncached)
            for i, idx in enumerate(uncached_idx):
                emb = new_embs[i]
                sent_emb[idx] = emb
                # Store in cache as bytes for memory efficiency
                cache_key = self._hash_text(uncached[i])
                self.emb_cache[cache_key] = emb.astype(np.float32).tobytes()

        return sent_emb


    def _regroup_embeddings(
        self,
        sent_emb: Dict[int, np.ndarray],
        sent_map: List[tuple[int, List[str]]],
    ) -> tuple[List[np.ndarray], List[List[str]]]:
        """
        Regroup flat embeddings back to per-prompt structure.

        Parameters
        ----------
        sent_emb : Dict[int, np.ndarray]
            Flat mapping: global_sentence_index → embedding.
        sent_map : List[Tuple[int, List[str]]]
            Mapping: (prompt_index, sentences_for_prompt).

        Returns
        -------
        Tuple[List[np.ndarray], List[List[str]]]
            - List of embedding arrays (one per prompt)
            - List of sentence lists (one per prompt)
        """
        emb_per_text: List[np.ndarray] = []
        sents_per_text: List[List[str]] = []

        offset = 0
        for prompt_idx, sents in sent_map:
            count = len(sents)
            # Gather embeddings for this prompt's sentences
            emb = np.array(
                [sent_emb[offset + j] for j in range(count)]
            )
            emb_per_text.append(emb)
            sents_per_text.append(sents)
            offset += count

        return emb_per_text, sents_per_text


    def _process_single(
        self,
        original: str,
        shielded: ShieldResult,
        sentences: List[str],
        embeddings: np.ndarray,
        aggressiveness: float,
        privacy_mode: bool,
        mode: str,
    ) -> Dict[str, Any]:
        """
        Process a single prompt through compression pipeline.

        This method executes stages 3-5 of the pipeline for one prompt:
        3. Compression: cluster sentences, select representatives
        4. Reconstruction: re-inject entities, compute metrics
        5. Safety: validate constraints, re-insert if needed

        Parameters
        ----------
        original : str
            Original prompt text (for token comparison).
        shielded : ShieldResult
            Result from shielding stage.
        sentences : List[str]
            Segmented sentences from shielded text.
        embeddings : np.ndarray
            Pre-computed embeddings for sentences.
        aggressiveness : float
            Compression intensity ∈ [0, 1].
        privacy_mode : bool
            Whether to suppress entity re-injection.
        mode : str
            Domain mode for safety validation.

        Returns
        -------
        Dict[str, Any]
            Dictionary with compressed text and metrics:
            - compressed_text: final output ready for LLM
            - original_tokens, compressed_tokens, tokens_saved
            - compression_ratio, cost_saved_usd
            - safety_score, alerts
        """
        # Stage 3: Semantic compression
        comp_sents, comp_stats = self.compressor.compress(
            sentences,
            embeddings,
            aggressiveness=aggressiveness,
            mode=mode,
        )
        comp_indices = comp_stats.get(
            "compressed_indices", list(range(len(comp_sents)))
        )

        # Stage 4: Reconstruction with token metrics
        recon: ReconstructionResult = self.reconstructor.reconstruct(
            original_prompt=original,
            compressed_sentences=comp_sents,
            shield_result=shielded,
            apply_stopwords=False,  # Stopwords handled in safety stage
            compressed_indices=comp_indices,
            privacy_mode=privacy_mode,
        )

        # Stage 5: Safety validation (may re-insert sentences)
        report: SafetyReport = self.safety.validate(
            original_text=original,
            compressed_text=recon.compressed_text,
            shield_result=shielded,
            original_sentences=sentences,
            compressed_indices=recon.compressed_indices,
            mode=mode,
        )

        # Start with safety-validated text
        final = report.final_text

        # Re-inject protected entities unless privacy mode is enabled
        if not privacy_mode:
            final = self.reconstructor._reinject_entities(
                final, shielded.placeholder_map
            )

        # Compute final token metrics
        final_tokens = len(self.reconstructor.tokenizer.encode(final))
        saved = recon.original_tokens - final_tokens

        # Deduplicate consecutive identical lines (post-processing)
        lines = final.splitlines()
        unique_lines: List[str] = []
        seen: set = set()
        for line in lines:
            stripped = line.strip()
            if stripped and stripped not in seen:
                unique_lines.append(stripped)
                seen.add(stripped)
        final = "\n".join(unique_lines) if unique_lines else final

        # Assemble result dictionary
        return {
            "compressed_text": final,
            "original_tokens": recon.original_tokens,
            "compressed_tokens": final_tokens,
            "tokens_saved": saved,
            "compression_ratio": (
                saved / recon.original_tokens if recon.original_tokens > 0 else 0.0
            ),
            "cost_saved_usd": saved * self.reconstructor.pricing["input"],
            "safety_score": report.safety_score,
            "alerts": recon.alerts + report.alerts,
            # Optional metadata for observability
            "_comp_stats": comp_stats,
            "_safety_report": {
                "missing_intents": report.missing_intents,
                "forced_added": report.forced_sentences_added,
                "perplexity": report.perplexity,
            },
        }


    def compress_batch(
        self,
        texts: List[str],
        aggressiveness: Optional[float] = None,
        mode: str = "general",
        nli_active: bool = False,
        language: Optional[str] = None,
        privacy_mode: Optional[bool] = None,
    ) -> List[Dict[str, Any]]:
        """
        Compress a batch of prompts through the full pipeline.

        This method orchestrates parallel execution of all pipeline stages
        for multiple prompts, maximizing throughput via thread pool.

        Parameters
        ----------
        texts : List[str]
            List of prompts to compress.
        aggressiveness : Optional[float], optional
            Compression intensity ∈ [0, 1]. If None, uses mode preset.
        mode : str, optional
            Domain mode: "legal", "finance", "code", or "general".
        nli_active : bool, optional
            Whether NLI-based validation is enabled (affects aggressiveness).
        language : Optional[str], optional
            Language code for sentence segmentation (auto-detect if None).
        privacy_mode : Optional[bool], optional
            Override instance default for PII handling.

        Returns
        -------
        List[Dict[str, Any]]
            List of result dictionaries, one per input prompt.
            Each contains compressed text and metrics (see _process_single).

        Pipeline Execution
        ------------------
        Stage 1: Shielding (parallel)
            - Check cache first, compute fresh if miss
            - Extract entities, restrictions, placeholders

        Stage 2: Segmentation + Embedding (batched)
            - Collect all sentences across prompts
            - Encode with caching to avoid redundant computation
            - Regroup embeddings by original prompt

        Stage 3-5: Compression + Reconstruction + Safety (parallel)
            - Process each prompt independently in thread pool
            - Apply domain-specific aggressiveness and constraints

        Complexity
        ----------
        Time: O(N · T_pipeline / M) amortized
            where N = prompt count, M = worker count,
            T_pipeline = latency for single prompt through all stages
        
        Space: O(N · L · d) for embeddings, L = avg sentences/prompt,
               d = embedding dimension (384 for MiniLM)

        Example
        -------
        >>> service = CompressionService(use_cache=True)
        >>> results = service.compress_batch(
        ...     texts=["Prompt 1", "Prompt 2"],
        ...     mode="code",
        ...     aggressiveness=0.3
        ... )
        >>> for i, r in enumerate(results):
        ...     print(f"Prompt {i}: saved {r['tokens_saved']} tokens")
        """
        if not texts:
            return []

        # Resolve effective aggressiveness
        effective_agg = self._compute_effective_aggressiveness(
            aggressiveness, mode, nli_active
        )
        logger.debug(f"Effective aggressiveness: {effective_agg:.2f} (mode={mode})")

        # Use instance privacy_mode if not overridden
        effective_privacy = (
            privacy_mode if privacy_mode is not None else self.privacy_mode
        )

        start_time = time.time()

        # Stage 1: Shielding (parallel with cache)
        shield_results: List[ShieldResult] = [None] * len(texts)  # type: ignore

        future_to_idx = {
            self.executor.submit(
                self._shield_with_cache, text, None, mode, effective_privacy
            ): i
            for i, text in enumerate(texts)
        }

        # Collect results as they complete
        for future in as_completed(future_to_idx):
            idx = future_to_idx[future]
            shield_results[idx] = future.result()

        # Stage 2: Segmentation + Embedding (batched for efficiency)
        all_sents, sent_map = self._collect_all_sentences(shield_results, language)

        if not all_sents:
            # Edge case: all prompts empty after shielding
            return [
                {
                    "compressed_text": "",
                    "original_tokens": 0,
                    "compressed_tokens": 0,
                    "tokens_saved": 0,
                    "compression_ratio": 0.0,
                    "cost_saved_usd": 0.0,
                    "safety_score": 1.0,
                    "alerts": ["Empty prompt after shielding"],
                }
                for _ in texts
            ]

        sent_emb = self._encode_with_cache(all_sents)
        emb_per_text, sents_per_text = self._regroup_embeddings(sent_emb, sent_map)

        # Stage 3-5: Compression + Reconstruction + Safety (parallel)
        results: List[Dict[str, Any]] = [None] * len(texts)  # type: ignore

        future_to_idx = {
            self.executor.submit(
                self._process_single,
                texts[i],
                shield_results[i],
                sents_per_text[i],
                emb_per_text[i],
                effective_agg,
                effective_privacy,
                mode,
            ): i
            for i in range(len(texts))
        }

        # Collect results as they complete
        for future in as_completed(future_to_idx):
            idx = future_to_idx[future]
            results[idx] = future.result()

        elapsed = time.time() - start_time
        logger.info(
            f"Batch compression complete: {len(texts)} prompts in {elapsed:.2f}s "
            f"({elapsed / len(texts) * 1000:.1f}ms/prompt avg)"
        )

        return results


    async def compress_batch_async(
        self,
        texts: List[str],
        aggressiveness: Optional[float] = None,
        mode: str = "general",
        nli_active: bool = False,
        language: Optional[str] = None,
        privacy_mode: Optional[bool] = None,
    ) -> List[Dict[str, Any]]:
        """
        Async wrapper for compress_batch (non-blocking event loop).

        Offloads CPU-bound compression to worker threads via asyncio.to_thread,
        preventing event loop starvation in async applications (FastAPI, etc.).

        Parameters
        ----------
        texts : List[str]
            List of prompts to compress.
        aggressiveness : Optional[float], optional
            Compression intensity ∈ [0, 1].
        mode : str, optional
            Domain mode for pattern selection.
        nli_active : bool, optional
            Whether NLI-based validation is enabled.
        language : Optional[str], optional
            Language code for segmentation.
        privacy_mode : Optional[bool], optional
            Override for PII handling.

        Returns
        -------
        List[Dict[str, Any]]
            Same as compress_batch().

        Note
        ----
        This does not provide true parallelism; it uses a thread pool
        to avoid blocking the async event loop. For true parallelism,
        consider multiprocessing or distributed processing.
        """
        return await asyncio.to_thread(
            self.compress_batch,
            texts,
            aggressiveness,
            mode,
            nli_active,
            language,
            privacy_mode,
        )


    def clear_caches(self, semantic_only: bool = False) -> Dict[str, int]:
        """
        Clear internal caches to free memory or force refresh.

        Parameters
        ----------
        semantic_only : bool, optional
            If True, only clear semantic cache (Redis).
            If False, clear all caches: shield, embeddings, semantic.

        Returns
        -------
        Dict[str, int]
            Count of cleared entries per cache type.
        """
        cleared: Dict[str, int] = {}

        if not semantic_only:
            # Clear in-memory shield cache
            if self.shield_cache is not None:
                cleared["shield_cache"] = len(self.shield_cache)
                self.shield_cache.clear()

            # Clear in-memory embedding cache
            if self.emb_cache is not None:
                cleared["emb_cache"] = len(self.emb_cache)
                self.emb_cache.clear()

        # Clear semantic cache (Redis)
        if self.semantic_cache:
            cleared["semantic_cache"] = self.semantic_cache.clear()

        logger.info(f"Caches cleared: {cleared}")
        return cleared


    def get_stats(self) -> Dict[str, Any]:
        """
        Return operational statistics for monitoring.

        Returns
        -------
        Dict[str, Any]
            Statistics including:
            - cache_sizes: counts for shield/emb caches
            - semantic_cache: stats from RedisVL index (if available)
            - configuration: current settings for observability
        """
        stats: Dict[str, Any] = {
            "cache_sizes": {
                "shield": len(self.shield_cache) if self.shield_cache else 0,
                "embeddings": len(self.emb_cache) if self.emb_cache else 0,
            },
            "configuration": {
                "use_cache": self.use_emb_cache,
                "privacy_mode": self.privacy_mode,
                "nli_enabled": self.nli_refinement_fn is not None,
            },
        }

        if self.semantic_cache:
            stats["semantic_cache"] = self.semantic_cache.get_stats()

        return stats