"""
Extraction Cache

Caches extracted state (intents, decisions) per user to enable incremental 
extraction - only processing new messages instead of re-processing entire 
conversation history.

This is a key optimization (Issue #9) that:
- Reduces latency for returning users by ~38%
- Avoids redundant LLM calls for already-processed messages
- Enables efficient multi-session conversations

Cache Strategy:
- Key: (user_id, messages_hash) where hash = hash of sorted message contents
- Value: ExtractionState containing last processed index + extracted data
- Invalidation: On explicit clear or when message history changes unexpectedly
"""

from __future__ import annotations

import hashlib
import json
import logging
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional
from datetime import datetime, timezone


def _utc_now() -> datetime:
    """Return current UTC time (timezone-aware)."""
    return datetime.now(timezone.utc)


logger = logging.getLogger(__name__)


@dataclass
class CacheStats:
    """Statistics for cache performance monitoring."""
    hits: int = 0
    misses: int = 0
    partial_hits: int = 0  # Had cache but new messages to process
    invalidations: int = 0
    
    @property
    def hit_rate(self) -> float:
        """Calculate cache hit rate."""
        total = self.hits + self.misses + self.partial_hits
        if total == 0:
            return 0.0
        # Count partial hits as half a hit for rate calculation
        return (self.hits + 0.5 * self.partial_hits) / total
    
    @property
    def total_requests(self) -> int:
        return self.hits + self.misses + self.partial_hits
    
    def to_dict(self) -> dict:
        return {
            "hits": self.hits,
            "misses": self.misses,
            "partial_hits": self.partial_hits,
            "invalidations": self.invalidations,
            "hit_rate": round(self.hit_rate, 3),
            "total_requests": self.total_requests
        }


@dataclass
class ExtractionState:
    """
    Cached extraction state for a user's conversation.
    
    Stores the extracted intents/decisions up to a certain point,
    allowing incremental extraction of only new messages.
    """
    user_id: str
    
    # Tracking what messages were processed
    messages_processed: int = 0  # Count of messages already extracted
    messages_hash: str = ""  # Hash of processed messages for validation
    
    # Extracted state (serialized for JSON storage)
    intent_graph_data: Optional[dict] = None
    decisions_data: list = field(default_factory=list)
    
    # Issue #118: Unified graph data (serialized UnifiedGraph)
    unified_graph_data: Optional[dict] = None
    
    # Issue #131: Track hashes of each processed message batch for per-turn callers.
    # Enables cache hits when callers pass only new messages (not full history).
    processed_hashes: list = field(default_factory=list)
    
    # Metadata
    created_at: datetime = field(default_factory=_utc_now)
    updated_at: datetime = field(default_factory=_utc_now)
    extraction_time_ms: int = 0  # Total time spent extracting
    
    def to_dict(self) -> dict:
        """Serialize to JSON-compatible dict."""
        return {
            "user_id": self.user_id,
            "messages_processed": self.messages_processed,
            "messages_hash": self.messages_hash,
            "intent_graph_data": self.intent_graph_data,
            "decisions_data": self.decisions_data,
            "unified_graph_data": self.unified_graph_data,
            "processed_hashes": self.processed_hashes,
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
            "extraction_time_ms": self.extraction_time_ms
        }
    
    @classmethod
    def from_dict(cls, data: dict) -> "ExtractionState":
        """Deserialize from dict."""
        return cls(
            user_id=data["user_id"],
            messages_processed=data.get("messages_processed", 0),
            messages_hash=data.get("messages_hash", ""),
            intent_graph_data=data.get("intent_graph_data"),
            decisions_data=data.get("decisions_data", []),
            unified_graph_data=data.get("unified_graph_data"),
            processed_hashes=data.get("processed_hashes", []),
            created_at=datetime.fromisoformat(data["created_at"]) if data.get("created_at") else _utc_now(),
            updated_at=datetime.fromisoformat(data["updated_at"]) if data.get("updated_at") else _utc_now(),
            extraction_time_ms=data.get("extraction_time_ms", 0)
        )


def compute_messages_hash(messages: list[dict], start_idx: int = 0, end_idx: Optional[int] = None) -> str:
    """
    Compute a stable hash of messages for cache validation.
    
    Uses content and role to create a deterministic hash that can detect
    if messages have been modified or reordered.
    
    Args:
        messages: List of message dicts with 'role' and 'content'
        start_idx: Starting index (inclusive)
        end_idx: Ending index (exclusive), None for all remaining
        
    Returns:
        SHA-256 hash string (first 16 chars for efficiency)
    """
    subset = messages[start_idx:end_idx] if end_idx else messages[start_idx:]
    
    # Create deterministic string from messages
    parts = []
    for msg in subset:
        role = msg.get("role", "unknown")
        content = msg.get("content", "")
        parts.append(f"{role}:{content}")
    
    combined = "|".join(parts)
    hash_obj = hashlib.sha256(combined.encode("utf-8"))
    return hash_obj.hexdigest()[:16]


class ExtractionCache:
    """
    In-memory cache for extraction state with optional disk persistence.
    
    Provides fast lookups for returning users and supports incremental
    extraction by tracking which messages have been processed.
    """
    
    def __init__(
        self, 
        persist_path: Optional[Path] = None,
        max_entries: int = 1000,
        ttl_seconds: int = 86400  # 24 hours default
    ):
        """
        Initialize extraction cache.
        
        Args:
            persist_path: Optional path for disk persistence
            max_entries: Maximum cache entries (LRU eviction)
            ttl_seconds: Time-to-live for cache entries
        """
        self._cache: dict[str, ExtractionState] = {}
        self._access_times: dict[str, float] = {}  # For LRU
        self._persist_path = persist_path
        self._max_entries = max_entries
        self._ttl_seconds = ttl_seconds
        self._stats = CacheStats()
        
        # Load from disk if persistence enabled
        if persist_path:
            self._load_from_disk()
    
    def get_state(
        self, 
        user_id: str, 
        messages: list[dict]
    ) -> tuple[Optional[ExtractionState], list[dict]]:
        """
        Get cached state and determine which messages need processing.
        
        Returns:
            Tuple of (cached_state, new_messages_to_process)
            - If full cache hit: (state, [])
            - If partial hit: (state, new_messages)
            - If miss: (None, all_messages)
        """
        cache_key = user_id
        
        # Check if we have cached state
        if cache_key not in self._cache:
            self._stats.misses += 1
            logger.debug(f"Cache MISS for user {user_id}")
            return None, messages
        
        state = self._cache[cache_key]
        self._access_times[cache_key] = time.time()
        
        # Check TTL
        age = (_utc_now() - state.updated_at).total_seconds()
        if age > self._ttl_seconds:
            self._stats.invalidations += 1
            logger.debug(f"Cache EXPIRED for user {user_id} (age: {age:.0f}s)")
            self.invalidate(user_id)
            return None, messages
        
        # Validate cached messages haven't changed
        if state.messages_processed > 0:
            cached_hash = compute_messages_hash(messages, 0, state.messages_processed)
            if cached_hash != state.messages_hash:
                # Prefix doesn't match — check if this is a per-turn caller
                # (Issue #131: real-world agents pass only new messages each turn,
                # not the full accumulated history)
                batch_hash = compute_messages_hash(messages)
                
                if batch_hash in state.processed_hashes:
                    # Same message batch already extracted → full hit (replay/duplicate)
                    self._stats.hits += 1
                    logger.info(
                        f"Cache HIT for user {user_id}: "
                        f"message batch already processed (per-turn pattern)"
                    )
                    return state, []
                
                if state.unified_graph_data:
                    # Have a cached graph — treat as partial hit (new turn to extract)
                    self._stats.partial_hits += 1
                    logger.info(
                        f"Cache PARTIAL HIT for user {user_id}: "
                        f"new message batch ({len(messages)} msgs), "
                        f"using cached graph as base (per-turn pattern)"
                    )
                    return state, messages
                
                # No usable graph — invalidate and reprocess from scratch
                self._stats.invalidations += 1
                logger.warning(
                    f"Cache INVALIDATED for user {user_id}: "
                    f"message history changed and no cached graph"
                )
                self.invalidate(user_id)
                return None, messages
        
        # Determine new messages (full-history caller pattern)
        new_messages = messages[state.messages_processed:]
        
        if not new_messages:
            # Full cache hit - no new messages
            self._stats.hits += 1
            logger.info(f"Cache HIT for user {user_id}: all {len(messages)} messages cached")
            return state, []
        else:
            # Partial hit - have cache but new messages
            self._stats.partial_hits += 1
            logger.info(
                f"Cache PARTIAL HIT for user {user_id}: "
                f"{state.messages_processed} cached, {len(new_messages)} new"
            )
            return state, new_messages
    
    def update_state(
        self,
        user_id: str,
        messages: list[dict],
        intent_graph_data: Optional[dict],
        decisions_data: list,
        extraction_time_ms: int = 0
    ) -> ExtractionState:
        """
        Update cached state after extraction.
        
        Args:
            user_id: User identifier
            messages: Full message list that was processed
            intent_graph_data: Serialized intent graph
            decisions_data: Serialized decisions list
            extraction_time_ms: Time taken for extraction
            
        Returns:
            Updated ExtractionState
        """
        cache_key = user_id
        
        # Get or create state
        if cache_key in self._cache:
            state = self._cache[cache_key]
            state.updated_at = _utc_now()
            state.extraction_time_ms += extraction_time_ms
        else:
            state = ExtractionState(user_id=user_id)
            state.extraction_time_ms = extraction_time_ms
        
        # Update with new data
        state.messages_processed = len(messages)
        state.messages_hash = compute_messages_hash(messages)
        state.intent_graph_data = intent_graph_data
        state.decisions_data = decisions_data
        
        # Issue #131: Track batch hash
        batch_hash = compute_messages_hash(messages)
        if batch_hash not in state.processed_hashes:
            state.processed_hashes.append(batch_hash)
            if len(state.processed_hashes) > 500:
                state.processed_hashes = state.processed_hashes[-500:]
        
        # Store in cache
        self._cache[cache_key] = state
        self._access_times[cache_key] = time.time()
        
        # Evict if over limit
        self._maybe_evict()
        
        # Persist if enabled
        if self._persist_path:
            self._save_to_disk()
        
        logger.debug(
            f"Cache UPDATED for user {user_id}: "
            f"{state.messages_processed} messages, "
            f"intent_graph={state.intent_graph_data is not None}, "
            f"decisions={len(state.decisions_data)}"
        )
        
        return state
    
    def update_unified_state(
        self,
        user_id: str,
        messages: list[dict],
        unified_graph_data: dict,
        extraction_time_ms: int = 0
    ) -> ExtractionState:
        """
        Update cached state with unified graph data after extraction (Issue #118).
        
        This is the primary cache update method for the unified extraction pipeline.
        Stores the serialized UnifiedGraph so subsequent add() calls can skip LLM
        extraction entirely (full hit) or extract only new messages (partial hit).
        
        Args:
            user_id: User identifier
            messages: Full message list that was processed
            unified_graph_data: Serialized UnifiedGraph (from graph.to_dict())
            extraction_time_ms: Time taken for extraction
            
        Returns:
            Updated ExtractionState
        """
        cache_key = user_id
        
        # Get or create state
        if cache_key in self._cache:
            state = self._cache[cache_key]
            state.updated_at = _utc_now()
            state.extraction_time_ms += extraction_time_ms
        else:
            state = ExtractionState(user_id=user_id)
            state.extraction_time_ms = extraction_time_ms
        
        # Update with unified graph data
        state.messages_processed = len(messages)
        state.messages_hash = compute_messages_hash(messages)
        state.unified_graph_data = unified_graph_data
        
        # Issue #131: Track this message batch hash for per-turn caller detection
        batch_hash = compute_messages_hash(messages)
        if batch_hash not in state.processed_hashes:
            state.processed_hashes.append(batch_hash)
            # Cap at 500 entries to bound memory usage
            if len(state.processed_hashes) > 500:
                state.processed_hashes = state.processed_hashes[-500:]
        
        # Store in cache
        self._cache[cache_key] = state
        self._access_times[cache_key] = time.time()
        
        # Evict if over limit
        self._maybe_evict()
        
        # Persist if enabled
        if self._persist_path:
            self._save_to_disk()
        
        logger.debug(
            f"Cache UPDATED (unified) for user {user_id}: "
            f"{state.messages_processed} messages, "
            f"graph_nodes={len(unified_graph_data.get('n', []))}"
        )
        
        return state
    
    def invalidate(self, user_id: str) -> bool:
        """
        Invalidate cache for a user.
        
        Args:
            user_id: User identifier
            
        Returns:
            True if entry was removed, False if not found
        """
        cache_key = user_id
        if cache_key in self._cache:
            del self._cache[cache_key]
            self._access_times.pop(cache_key, None)
            self._stats.invalidations += 1
            logger.info(f"Cache INVALIDATED for user {user_id}")
            
            if self._persist_path:
                self._save_to_disk()
            return True
        return False
    
    def clear(self) -> int:
        """
        Clear entire cache.
        
        Returns:
            Number of entries cleared
        """
        count = len(self._cache)
        self._cache.clear()
        self._access_times.clear()
        logger.info(f"Cache CLEARED: {count} entries removed")
        
        if self._persist_path:
            self._save_to_disk()
        
        return count
    
    def get_stats(self) -> CacheStats:
        """Get cache statistics."""
        return self._stats
    
    def reset_stats(self) -> None:
        """Reset cache statistics."""
        self._stats = CacheStats()
    
    def _maybe_evict(self) -> None:
        """Evict least recently used entries if over limit."""
        while len(self._cache) > self._max_entries:
            # Find LRU entry
            lru_key = min(self._access_times.keys(), key=lambda k: self._access_times[k])
            del self._cache[lru_key]
            del self._access_times[lru_key]
            logger.debug(f"Cache EVICTED user {lru_key} (LRU)")
    
    def _save_to_disk(self) -> None:
        """Save cache to disk."""
        if not self._persist_path:
            return
        
        try:
            self._persist_path.parent.mkdir(parents=True, exist_ok=True)
            data = {
                "entries": {k: v.to_dict() for k, v in self._cache.items()},
                "stats": self._stats.to_dict(),
                "saved_at": _utc_now().isoformat()
            }
            with open(self._persist_path, 'w') as f:
                json.dump(data, f, indent=2)
            logger.debug(f"Cache saved to {self._persist_path}")
        except Exception as e:
            logger.error(f"Failed to save cache: {e}")
    
    def _load_from_disk(self) -> None:
        """Load cache from disk."""
        if not self._persist_path or not self._persist_path.exists():
            return
        
        try:
            with open(self._persist_path, 'r') as f:
                data = json.load(f)
            
            for key, entry_data in data.get("entries", {}).items():
                state = ExtractionState.from_dict(entry_data)
                # Check TTL on load
                age = (_utc_now() - state.updated_at).total_seconds()
                if age <= self._ttl_seconds:
                    self._cache[key] = state
                    self._access_times[key] = time.time()
            
            logger.info(f"Cache loaded from {self._persist_path}: {len(self._cache)} entries")
        except Exception as e:
            logger.error(f"Failed to load cache: {e}")
