"""
Core ContextCompressor implementation - the main entry point for the compression layer.
"""

import time
from typing import Optional, Any, Dict
from dataclasses import dataclass

import torch

from ares_compression.config import CompressionConfig, CompressionMode
from ares_compression.layers import (
    SemanticTokenPruning,
    DynamicPrecisionSwitching,
    EntropyGatedCompression,
    TieredCacheManager,
)
from ares_compression.lobes import MultiLobeSystem


@dataclass
class CompressionResult:
    """Result of context compression."""
    compressed_context: torch.Tensor
    attention_mask: Optional[torch.Tensor]
    original_length: int
    compressed_length: int
    compression_ratio: float
    memory_saved: float  # in MB
    latency_ms: float
    metadata: Dict[str, Any]


def _resolve_device(device: Optional[str]) -> str:
    """Resolve the requested execution device against local hardware."""
    requested = (device or "").strip().lower()
    if not requested:
        requested = "cuda" if torch.cuda.is_available() else "cpu"
    if requested.startswith("cuda") and not torch.cuda.is_available():
        return "cpu"
    try:
        torch.device(requested)
    except (RuntimeError, ValueError):
        return "cpu"
    return requested


class ContextCompressor:
    """
    Main context compression layer.

    This is the primary interface for compressing LLM contexts.
    It orchestrates multiple compression strategies:
    - Semantic token pruning
    - Dynamic precision switching
    - Entropy-gated compression
    - Tiered caching

    Example:
        >>> compressor = ContextCompressor.from_config(CompressionConfig())
        >>> result = compressor.compress(context_tensor)
        >>> print(f"Compressed {result.original_length} -> {result.compressed_length}")
    """

    def __init__(
        self,
        config: Optional[CompressionConfig] = None,
        device: Optional[str] = None
    ):
        """
        Initialize the compressor.

        Args:
            config: Compression configuration (uses BALANCED mode if None)
            device: Device to run on ("cuda" or "cpu")
        """
        self.config = config or CompressionConfig(mode=CompressionMode.BALANCED)
        self.device = _resolve_device(device)

        # Initialize layers based on config
        self.semantic_pruner: Optional[SemanticTokenPruning] = None
        self.precision_switcher: Optional[DynamicPrecisionSwitching] = None
        self.entropy_gater: Optional[EntropyGatedCompression] = None
        self.cache_manager: Optional[TieredCacheManager] = None

        if self.config.semantic_pruning and self.config.semantic_pruning.enabled:
            self.semantic_pruner = SemanticTokenPruning(
                self.config.semantic_pruning, self.device
            )

        if self.config.dynamic_precision and self.config.dynamic_precision.enabled:
            self.precision_switcher = DynamicPrecisionSwitching(
                self.config.dynamic_precision, self.device
            )

        if self.config.entropy_gating and self.config.entropy_gating.enabled:
            self.entropy_gater = EntropyGatedCompression(
                self.config.entropy_gating, self.device
            )

        if self.config.tiered_cache and self.config.tiered_cache.enabled:
            self.cache_manager = TieredCacheManager(
                self.config.tiered_cache, self.device
            )

        # Step 0: Multi-Lobe Router (MAS-RON Core)
        self.multi_lobe: Optional[MultiLobeSystem] = None
        # d_model will be initialized on first compress call if needed

        # Statistics (Restored)
        self.stats = {
            "total_compressions": 0,
            "total_memory_saved_mb": 0.0,
            "total_latency_ms": 0.0,
            "avg_compression_ratio": 0.0,
        }

    def _init_multi_lobe(self, d_model: int):
        self.multi_lobe = MultiLobeSystem(d_model, self.device).to(self.device)

    def compress(
        self,
        context: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        **kwargs
    ) -> CompressionResult:
        """
        Compress a context tensor.

        Args:
            context: Input context tensor [seq_len, hidden_dim] or [batch, seq_len, hidden_dim]
            attention_mask: Optional attention mask [seq_len] or [batch, seq_len]
            **kwargs: Additional arguments for compression layers

        Returns:
            CompressionResult with compressed context and metadata
        """
        start_time = time.time()

        # Ensure inputs are on the active device.
        context = context.to(self.device)
        if attention_mask is not None:
            attention_mask = attention_mask.to(self.device)
        
        # Lazy init for multi-lobe based on d_model
        if self.multi_lobe is None:
            self._init_multi_lobe(context.shape[-1])

        original_shape = context.shape
        original_length = original_shape[-2]

        # Add batch dimension if needed
        squeeze_output = False
        if context.dim() == 2:
            context = context.unsqueeze(0)
            squeeze_output = True
            if attention_mask is not None and attention_mask.dim() == 1:
                attention_mask = attention_mask.unsqueeze(0)

        # Store original memory size
        original_memory_mb = context.element_size() * context.numel() / (1024 * 1024)

        # Step 0: Apply Multi-Lobe Routing (MAS-RON Enhancement)
        # This enriches the context using the specialist lobes (General/Code)
        enriched_context = self.multi_lobe(context)

        # Apply compression pipeline to the ENRICHED context
        compressed_context = enriched_context
        compressed_attention_mask = attention_mask
        metadata = {}

        # Step 1: Semantic Token Pruning
        if self.semantic_pruner:
            compressed_context, compressed_attention_mask, prune_metadata = self.semantic_pruner.prune(
                compressed_context, compressed_attention_mask
            )
            metadata["pruning"] = prune_metadata

        # Step 2: Entropy-Gated Compression
        if self.entropy_gater:
            compressed_context, entropy_metadata = self.entropy_gater.compress(
                compressed_context
            )
            metadata["entropy"] = entropy_metadata

        # Step 3: Dynamic Precision Switching
        if self.precision_switcher:
            compressed_context, precision_metadata = self.precision_switcher.compress(
                compressed_context
            )
            metadata["precision"] = precision_metadata

        # Step 4: Cache Management
        if self.cache_manager:
            self.cache_manager.update(compressed_context)

        # Remove batch dimension if added
        if squeeze_output:
            compressed_context = compressed_context.squeeze(0)
            if compressed_attention_mask is not None:
                compressed_attention_mask = compressed_attention_mask.squeeze(0)

        # Calculate metrics
        compressed_length = compressed_context.shape[-2]
        compressed_memory_mb = (
            compressed_context.element_size() * compressed_context.numel() / (1024 * 1024)
        )

        compression_ratio = 1.0 - (compressed_length / original_length)
        memory_saved = original_memory_mb - compressed_memory_mb
        latency_ms = (time.time() - start_time) * 1000

        # Update statistics
        self.stats["total_compressions"] += 1
        self.stats["total_memory_saved_mb"] += memory_saved
        self.stats["total_latency_ms"] += latency_ms
        self.stats["avg_compression_ratio"] = (
            (self.stats["avg_compression_ratio"] * (self.stats["total_compressions"] - 1) + compression_ratio)
            / self.stats["total_compressions"]
        )

        return CompressionResult(
            compressed_context=compressed_context,
            attention_mask=compressed_attention_mask,
            original_length=original_length,
            compressed_length=compressed_length,
            compression_ratio=compression_ratio,
            memory_saved=memory_saved,
            latency_ms=latency_ms,
            metadata=metadata
        )

    def get_stats(self) -> Dict[str, Any]:
        """Get compression statistics."""
        return self.stats.copy()

    def reset_stats(self):
        """Reset compression statistics."""
        self.stats = {
            "total_compressions": 0,
            "total_memory_saved_mb": 0.0,
            "total_latency_ms": 0.0,
            "avg_compression_ratio": 0.0,
        }

    def to(self, device: str) -> "ContextCompressor":
        """Move compressor to different device."""
        self.device = _resolve_device(device)
        if self.semantic_pruner:
            self.semantic_pruner.to(self.device)
        if self.precision_switcher:
            self.precision_switcher.to(self.device)
        if self.entropy_gater:
            self.entropy_gater.to(self.device)
        if self.cache_manager:
            self.cache_manager.to(self.device)
        return self

    def __call__(self, *args, **kwargs) -> CompressionResult:
        """Allow calling compressor directly."""
        return self.compress(*args, **kwargs)
