"""
Configuration models for the ARES Context Compression Layer.
"""

from enum import Enum
from typing import Literal, Optional
from pydantic import BaseModel, ConfigDict, Field


class CompressionMode(str, Enum):
    """Compression mode presets."""
    AGGRESSIVE = "aggressive"
    BALANCED = "balanced"
    CONSERVATIVE = "conservative"
    CUSTOM = "custom"


class SemanticPruningConfig(BaseModel):
    """Configuration for semantic token pruning."""
    enabled: bool = True
    target_ratio: float = Field(
        default=0.4,
        ge=0.0,
        le=1.0,
        description="Target ratio of tokens to keep (0.4 = keep 40%)"
    )
    min_tokens: int = Field(
        default=128,
        ge=0,
        description="Minimum number of tokens to preserve"
    )
    use_attention_scores: bool = True
    use_position_aware: bool = True


class DynamicPrecisionConfig(BaseModel):
    """Configuration for dynamic precision switching."""
    enabled: bool = True
    default_precision: Literal["fp16", "int8", "int4"] = "fp16"
    precision_thresholds: dict[str, float] = Field(
        default={
            "fp16": 0.8,
            "int8": 0.5,
            "int4": 0.2
        },
        description="Entropy thresholds for precision switching"
    )
    switch_latency_penalty: float = Field(
        default=0.001,
        ge=0.0,
        description="Latency penalty for precision switches"
    )


class EntropyGatingConfig(BaseModel):
    """Configuration for entropy-gated compression."""
    enabled: bool = True
    low_entropy_threshold: float = Field(
        default=0.3,
        ge=0.0,
        le=1.0,
        description="Entropy threshold for low-entropy tokens"
    )
    high_entropy_threshold: float = Field(
        default=0.7,
        ge=0.0,
        le=1.0,
        description="Entropy threshold for high-entropy tokens"
    )


class TieredCacheConfig(BaseModel):
    """Configuration for tiered caching."""
    enabled: bool = True
    hot_cache_size: int = Field(
        default=1000,
        ge=0,
        description="Number of states in hot cache (FP16)"
    )
    warm_cache_size: int = Field(
        default=5000,
        ge=0,
        description="Number of states in warm cache (INT8)"
    )
    cold_cache_size: int = Field(
        default=10000,
        ge=0,
        description="Number of states in cold cache (INT4/compressed)"
    )
    eviction_policy: Literal["lru", "lfu", "fifo"] = "lru"


class CompressionConfig(BaseModel):
    """Master configuration for the compression layer."""
    model_config = ConfigDict(arbitrary_types_allowed=True)

    mode: CompressionMode = CompressionMode.BALANCED
    semantic_pruning: Optional[SemanticPruningConfig] = None
    dynamic_precision: Optional[DynamicPrecisionConfig] = None
    entropy_gating: Optional[EntropyGatingConfig] = None
    tiered_cache: Optional[TieredCacheConfig] = None

    # Performance targets
    max_memory_reduction: float = Field(
        default=0.5,
        ge=0.0,
        le=1.0,
        description="Target memory reduction (0.5 = 50%)"
    )
    max_latency_overhead: float = Field(
        default=0.1,
        ge=0.0,
        description="Maximum acceptable latency overhead (10%)"
    )

    def __init__(self, **data):
        # Handle mode as string or enum
        if 'mode' in data and isinstance(data['mode'], str):
            data['mode'] = CompressionMode(data['mode'])
        super().__init__(**data)
        # Apply mode presets if mode is not custom
        if self.mode != CompressionMode.CUSTOM:
            self._apply_mode_preset()

    def _apply_mode_preset(self):
        """Apply preset configurations based on mode."""
        if self.mode == CompressionMode.AGGRESSIVE:
            self.semantic_pruning = SemanticPruningConfig(
                enabled=True, target_ratio=0.3, min_tokens=64
            )
            self.dynamic_precision = DynamicPrecisionConfig(
                enabled=True,
                default_precision="int8",
                precision_thresholds={"fp16": 0.9, "int8": 0.6, "int4": 0.3}
            )
            self.entropy_gating = EntropyGatingConfig(
                enabled=True, low_entropy_threshold=0.4
            )
            self.tiered_cache = TieredCacheConfig(
                enabled=True,
                hot_cache_size=500,
                warm_cache_size=3000,
                cold_cache_size=8000
            )
            self.max_memory_reduction = 0.6

        elif self.mode == CompressionMode.BALANCED:
            self.semantic_pruning = SemanticPruningConfig(
                enabled=True, target_ratio=0.4, min_tokens=128
            )
            self.dynamic_precision = DynamicPrecisionConfig(
                enabled=True,
                default_precision="fp16",
                precision_thresholds={"fp16": 0.8, "int8": 0.5, "int4": 0.2}
            )
            self.entropy_gating = EntropyGatingConfig(
                enabled=True, low_entropy_threshold=0.3, high_entropy_threshold=0.7
            )
            self.tiered_cache = TieredCacheConfig(
                enabled=True,
                hot_cache_size=1000,
                warm_cache_size=5000,
                cold_cache_size=10000
            )
            self.max_memory_reduction = 0.5

        elif self.mode == CompressionMode.CONSERVATIVE:
            self.semantic_pruning = SemanticPruningConfig(
                enabled=True, target_ratio=0.6, min_tokens=256
            )
            self.dynamic_precision = DynamicPrecisionConfig(
                enabled=True,
                default_precision="fp16",
                precision_thresholds={"fp16": 0.7, "int8": 0.4}
            )
            self.entropy_gating = EntropyGatingConfig(
                enabled=True, low_entropy_threshold=0.2, high_entropy_threshold=0.8
            )
            self.tiered_cache = TieredCacheConfig(
                enabled=True,
                hot_cache_size=2000,
                warm_cache_size=8000,
                cold_cache_size=15000
            )
            self.max_memory_reduction = 0.4
