"""Cache system for evaluation results."""

import sqlite3
import pickle
import hashlib
import json
import os
from pathlib import Path
from typing import Any, Optional, Dict
from datetime import datetime, timedelta



# Helpers
def get_default_cache_dir() -> Path:
    """Get platform-specific default cache directory."""
    if os.name == 'nt':  # Windows
        cache_home = Path(os.getenv('LOCALAPPDATA', os.path.expanduser('~')))
    else:  # Linux/Mac
        cache_home = Path(os.getenv('XDG_CACHE_HOME', os.path.expanduser('~/.cache')))
    
    return cache_home / 'openevalkit'

def make_cache_key(run, evaluator) -> str:
    """
    Create cache key from run and evaluator.
    
    Key includes everything that affects the result:
    - Run data (input, output, reference)
    - Evaluator identity and configuration
    
    Args:
        run: Run object
        evaluator: Scorer or Judge instance
        
    Returns:
        SHA256 hash (16 chars) as cache key
        
    Examples:
        >>> from openevalkit import Run
        >>> from openevalkit.scorers import ExactMatch
        >>> run = Run(id="1", input="Q", output="A", reference="A")
        >>> scorer = ExactMatch()
        >>> key = make_cache_key(run, scorer)
        >>> len(key)
        16
    """
    import openevalkit
    key_data = {
        "library_version": openevalkit.__version__,
        "input": str(run.input),
        "output": str(run.output),
        "reference": str(run.reference) if run.reference else None,
        "evaluator_name": evaluator.name,
        "evaluator_config": _get_evaluator_config_hash(evaluator)
    }
    
    key_str = json.dumps(key_data, sort_keys=True)
    return hashlib.sha256(key_str.encode()).hexdigest()[:16]



def _get_evaluator_config_hash(evaluator) -> str:
    """
    Get hash of evaluator configuration that affects results.
    
    Args:
        evaluator: Scorer or Judge instance
        
    Returns:
        Config hash string
    """
    from openevalkit.judges.base import Judge
    
    config = {"type": "judge" if isinstance(evaluator, Judge) else "scorer"}
    
    if isinstance(evaluator, Judge):
        # For judges: include LLM config and prompt
        if hasattr(evaluator, 'llm_config'):
            config.update({
                "model": evaluator.llm_config.model,
                "temperature": evaluator.llm_config.temperature,
            })
        
        if hasattr(evaluator, 'prompt_template'):
            config["prompt_hash"] = evaluator.prompt_template.hash()
        
        if hasattr(evaluator, 'rubric'):
            config.update({
                "criteria": evaluator.rubric.criteria,
                "scale": evaluator.rubric.scale,
                "weights": evaluator.rubric.weights
            })
        
        # For ensemble judges, include sub-judges
        if hasattr(evaluator, 'judges'):
            config["sub_judges"] = [
                _get_evaluator_config_hash(j) for j in evaluator.judges
            ]
            config["method"] = evaluator.method
    
    return hashlib.sha256(
        json.dumps(config, sort_keys=True).encode()
    ).hexdigest()[:16]



class Cache:
    """
    SQLite-based cache for evaluation results.
    
    Caches Score and Judgment objects to avoid recomputation.
    Implements LRU eviction and age-based cleanup.
    
    Examples:
        >>> cache = Cache()
        >>> cache.set("key1", {"score": 0.9})
        >>> cache.get("key1")
        {'score': 0.9}
        
        >>> # Custom location and limits
        >>> cache = Cache(
        ...     cache_dir="/shared/cache",
        ...     max_size_mb=1000
        ... )
    """
    
    def __init__(
        self,
        cache_dir: Optional[str] = None,
        max_size_mb: int = 500
    ):
        """
        Initialize cache with SQLite backend.
        
        Args:
            cache_dir: Directory for cache database (None = default)
            max_size_mb: Maximum cache size in MB before LRU eviction
        """
        if cache_dir is None:
            self.cache_dir = get_default_cache_dir()
        else:
            self.cache_dir = Path(cache_dir).expanduser()
        
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self.db_path = self.cache_dir / "cache.db"
        self.max_size_bytes = max_size_mb * 1024 * 1024
        
        self._init_db()
    
    def _init_db(self):
        """Create cache table if not exists."""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS cache (
                    key TEXT PRIMARY KEY,
                    value BLOB NOT NULL,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    accessed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            # Index for LRU eviction (find least recently accessed)
            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_accessed_at 
                ON cache(accessed_at)
            """)
            # Index for age-based cleanup
            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_created_at 
                ON cache(created_at)
            """)
    
    def get(self, key: str) -> Optional[Any]:
        """
        Get cached value by key.
        
        Updates accessed_at timestamp for LRU tracking.
        
        Args:
            key: Cache key
            
        Returns:
            Cached value or None if not found
        """
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute(
                "SELECT value FROM cache WHERE key = ?",
                (key,)
            )
            row = cursor.fetchone()
            
            if row:
                # Update accessed_at for LRU
                conn.execute(
                    "UPDATE cache SET accessed_at = ? WHERE key = ?",
                    (datetime.now(), key)
                )
                return pickle.loads(row[0])
        
        return None
    

    def set(self, key: str, value: Any):
        """
        Store value in cache.
        
        Automatically evicts old entries if cache exceeds size limit.
        
        Args:
            key: Cache key
            value: Value to cache (must be picklable)
        """
        # Check if we need to evict BEFORE adding new entry
        value_bytes = pickle.dumps(value)
        current_size = self._get_size_bytes()
        
        # If adding this entry would exceed limit, evict first
        if current_size + len(value_bytes) > self.max_size_bytes:
            # Evict to 80% to give headroom
            target_size = int(self.max_size_bytes * 0.8)
            self._evict_lru(target_size=target_size)
        
        # Store new entry
        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                """
                INSERT OR REPLACE INTO cache (key, value, created_at, accessed_at) 
                VALUES (?, ?, ?, ?)
                """,
                (key, value_bytes, datetime.now(), datetime.now())
            )


    def _evict_lru(self, target_size: float):
        """
        Evict least recently used entries until under target size.
        
        Args:
            target_size: Target size in bytes to evict down to
        """
        with sqlite3.connect(self.db_path) as conn:
            # Get all entries ordered by access time (oldest first) with sizes
            cursor = conn.execute("""
                SELECT key, LENGTH(value) as size 
                FROM cache 
                ORDER BY accessed_at ASC
            """)
            
            entries = cursor.fetchall()
            
            # Calculate current total size
            current_size = sum(size for _, size in entries)
            
            # If already under target, nothing to do
            if current_size <= target_size:
                return
            
            # Delete entries (oldest first) until we're under target
            keys_to_delete = []
            for key, size in entries:
                if current_size <= target_size:
                    break
                keys_to_delete.append(key)
                current_size -= size
            
            # Execute batch delete
            if keys_to_delete:
                placeholders = ','.join('?' * len(keys_to_delete))
                conn.execute(
                    f"DELETE FROM cache WHERE key IN ({placeholders})",
                    keys_to_delete
                )


    def clear(self):
        """Clear all cached entries."""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("DELETE FROM cache")
    
    def clear_old(self, days: int = 30):
        """
        Clear entries older than N days.
        
        Args:
            days: Delete entries created more than this many days ago
        """
        cutoff = datetime.now() - timedelta(days=days)
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute(
                "DELETE FROM cache WHERE created_at < ?",
                (cutoff,)
            )
            deleted = cursor.rowcount
        
        return deleted
    
    def _get_size_bytes(self) -> int:
        """Get total cache size in bytes."""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute(
                "SELECT SUM(LENGTH(value)) FROM cache"
            )
            size = cursor.fetchone()[0]
        
        return size or 0
    
    def stats(self) -> Dict[str, Any]:
        """
        Get cache statistics.
        
        Returns:
            Dictionary with entries count, size in MB, and DB path
        """
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.execute("SELECT COUNT(*) FROM cache")
            count = cursor.fetchone()[0]
            
            size_bytes = self._get_size_bytes()
        
        return {
            "entries": count,
            "size_mb": size_bytes / (1024 * 1024),
            "db_path": str(self.db_path)
        }