"""
CAAE-Inspired Response Cache for the kvict Proxy.

Uses a cost-model scoring system adapted from packages/caae/cost_model.py
to decide which cached responses to retain under memory pressure.

The CAAE cost model philosophy:
  - In CAAE, expensive-to-recompute KV blocks get higher retention priority.
  - In the proxy, expensive-to-regenerate API responses get higher retention priority.

Scoring formula (per cache entry):
  retention_score = (prompt_tokens * W_tokens + model_cost * W_cost) * decay(age)

  - prompt_tokens: longer prompts cost more to re-send (more input tokens billed)
  - model_cost: expensive models (GPT-4, Claude Opus) cost more to regenerate
  - decay(age): exponential decay based on entry age, so stale items get evicted first

This mirrors CAAE's:
  score = (layer_depth * W_layer + seq_index * W_seq) * exp(-decay * elapsed_time)
  where layer_depth -> model_cost, seq_index -> prompt_tokens

Cache types:
  1. Exact match: SHA-256 of (model, messages, temperature). Full cache hit.
  2. Prefix match: (future) shared system prompts across conversations.
"""

from __future__ import annotations

import hashlib
import json
import math
import time
import logging
from dataclasses import dataclass, field
from typing import Any, Dict, Optional

from proxy.config import ProxyConfig, get_config
from proxy.schemas import ChatCompletionRequest, ChatCompletionResponse

logger = logging.getLogger(__name__)


@dataclass
class CacheEntry:
    """A cached API response with cost-model metadata."""

    key: str
    response: ChatCompletionResponse
    prompt_tokens: int
    completion_tokens: int
    model_input_cost_per_1k: float
    created_at: float = field(default_factory=time.time)
    last_accessed: float = field(default_factory=time.time)
    access_count: int = 0

    def retention_score(self, config: ProxyConfig, current_time: float) -> float:
        """Compute CAAE-inspired retention priority score.

        Higher score = more important to keep. Mirrors the CAAE cost model's
        approach of preserving expensive-to-recompute blocks.

        The formula adapts CostModel.estimate_recompute_cost:
          CAAE:  base_cost * layer_factor * seq_factor * exp(-decay * elapsed)
          Proxy: (tokens * W_t + cost * W_c) * exp(-decay * age_seconds)
        """
        # Base value: how expensive this entry is to regenerate
        token_value = self.prompt_tokens * config.cost_model_weight_prompt_tokens
        cost_value = (
            self.model_input_cost_per_1k
            * 1000  # Scale up for scoring
            * config.cost_model_weight_model_cost
        )
        base_score = token_value + cost_value

        # Time decay: older entries lose priority (like CAAE's elapsed-time decay)
        age_seconds = current_time - self.last_accessed
        decay = math.exp(-config.cost_model_decay_rate * age_seconds)

        # Access frequency bonus (popular entries retained longer; strong enough
        # to outweigh cost difference so hot cheap entries can survive cold expensive ones)
        frequency_bonus = 1.0 + math.log1p(self.access_count) * 3.5

        return base_score * decay * frequency_bonus


class ResponseCache:
    """In-memory response cache with CAAE cost-model eviction.

    When the cache exceeds max_entries, the entry with the LOWEST
    retention_score is evicted — meaning cheap, stale, rarely-accessed
    responses are removed first, while expensive/hot responses survive.
    """

    def __init__(self, config: ProxyConfig | None = None):
        self._config = config or get_config()
        self._entries: Dict[str, CacheEntry] = {}
        self._stats = CacheStats()

    @property
    def size(self) -> int:
        return len(self._entries)

    @property
    def stats(self) -> "CacheStats":
        return self._stats

    def get(self, request: ChatCompletionRequest) -> Optional[ChatCompletionResponse]:
        """Look up a cached response for the given request.

        Returns None on cache miss.
        """
        key = self._compute_key(request)
        entry = self._entries.get(key)

        if entry is None:
            self._stats.misses += 1
            return None

        # Check TTL expiry
        now = time.time()
        if now - entry.created_at > self._config.cache_ttl_seconds:
            del self._entries[key]
            self._stats.misses += 1
            self._stats.ttl_evictions += 1
            return None

        # Cache hit — update access metadata
        entry.last_accessed = now
        entry.access_count += 1
        self._stats.hits += 1
        return entry.response

    def put(
        self,
        request: ChatCompletionRequest,
        response: ChatCompletionResponse,
        prompt_tokens: int,
        completion_tokens: int,
        model_input_cost_per_1k: float,
    ) -> None:
        """Store a response in the cache.

        If the cache is full, evict the entry with the lowest CAAE retention score.
        """
        key = self._compute_key(request)

        # Evict if at capacity
        if (
            len(self._entries) >= self._config.cache_max_entries
            and key not in self._entries
        ):
            self._evict_lowest_score()

        self._entries[key] = CacheEntry(
            key=key,
            response=response,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            model_input_cost_per_1k=model_input_cost_per_1k,
            created_at=time.time(),
            last_accessed=time.time(),
            access_count=0,
        )

    def clear(self) -> None:
        """Clear the entire cache."""
        self._entries.clear()

    def _compute_key(self, request: ChatCompletionRequest) -> str:
        """Compute a deterministic cache key from the request.

        Uses SHA-256 of the serialized (model, messages, temperature) tuple.
        Streaming flag is excluded — same content should share cache.
        """
        messages_data = [
            {"role": m.role, "content": m.content, "name": m.name}
            for m in request.messages
        ]
        key_material = json.dumps(
            {
                "model": request.model,
                "messages": messages_data,
                "temperature": request.temperature,
            },
            sort_keys=True,
            ensure_ascii=True,
        )
        return hashlib.sha256(key_material.encode()).hexdigest()

    def _evict_lowest_score(self) -> None:
        """Evict the cache entry with the lowest CAAE retention score.

        This is the core CAAE integration: instead of simple LRU, we score
        each entry by how expensive it would be to regenerate (token count +
        model cost) weighted by recency. Cheap/stale entries are evicted first.

        Mirrors CAAE's eviction philosophy from kv_evict/policies/lwkcp.py:
        "Preserve blocks that are expensive to recompute."
        """
        if not self._entries:
            return

        now = time.time()
        worst_key = None
        worst_score = float("inf")

        for key, entry in self._entries.items():
            score = entry.retention_score(self._config, now)
            if score < worst_score:
                worst_score = score
                worst_key = key

        if worst_key is not None:
            del self._entries[worst_key]
            self._stats.cost_model_evictions += 1
            logger.debug(
                "CAAE eviction: removed entry with score %.2f (cache size: %d)",
                worst_score,
                len(self._entries),
            )


@dataclass
class CacheStats:
    """Cache performance statistics."""

    hits: int = 0
    misses: int = 0
    ttl_evictions: int = 0
    cost_model_evictions: int = 0

    @property
    def total_requests(self) -> int:
        return self.hits + self.misses

    @property
    def hit_rate(self) -> float:
        total = self.total_requests
        if total == 0:
            return 0.0
        return self.hits / total
