"""
Advanced caching and batch operations for SLTK API.

Provides performance optimizations for large dataset operations including:
- Content-addressable caching for segment predictions
- Metric computation caching with automatic invalidation
- Batch operations for segment updates and metric computation
- Lazy loading utilities for large segment collections
"""

from __future__ import annotations

import hashlib
import logging
import threading
import time
from collections import OrderedDict
from collections.abc import Callable
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Generic, TypeVar

logger = logging.getLogger(__name__)

T = TypeVar("T")
K = TypeVar("K")


# ============================================================================
# LRU Cache Implementation
# ============================================================================


class LRUCache(Generic[K, T]):
    """
    Thread-safe LRU cache with TTL support.

    Features:
    - Configurable max size and TTL
    - Thread-safe operations
    - Statistics tracking (hits, misses)
    - Optional callback on eviction

    Example:
        >>> cache = LRUCache[str, dict](max_size=100, ttl_seconds=300)
        >>> cache.set("key1", {"data": "value"})
        >>> result = cache.get("key1")
    """

    def __init__(
        self,
        max_size: int = 100,
        ttl_seconds: float = 300,
        on_evict: Callable[[K, T], None] | None = None,
    ):
        self._cache: OrderedDict[K, tuple[T, float]] = OrderedDict()
        self._lock = threading.RLock()
        self._max_size = max_size
        self._ttl = ttl_seconds
        self._on_evict = on_evict
        self._hits = 0
        self._misses = 0

    def get(self, key: K) -> T | None:
        """Get value from cache, returning None if not found or expired."""
        with self._lock:
            if key not in self._cache:
                self._misses += 1
                return None

            value, timestamp = self._cache[key]
            current_time = time.time()

            # Check TTL
            if current_time - timestamp > self._ttl:
                self._evict(key)
                self._misses += 1
                return None

            # Move to end (most recently used)
            self._cache.move_to_end(key)
            self._hits += 1
            return value

    def set(self, key: K, value: T) -> None:
        """Set value in cache."""
        with self._lock:
            current_time = time.time()

            # Update existing key
            if key in self._cache:
                self._cache[key] = (value, current_time)
                self._cache.move_to_end(key)
                return

            # Evict oldest if at capacity
            while len(self._cache) >= self._max_size:
                oldest_key = next(iter(self._cache))
                self._evict(oldest_key)

            self._cache[key] = (value, current_time)

    def delete(self, key: K) -> bool:
        """Delete key from cache, returns True if key existed."""
        with self._lock:
            if key in self._cache:
                self._evict(key)
                return True
            return False

    def clear(self) -> None:
        """Clear all cache entries."""
        with self._lock:
            self._cache.clear()

    def _evict(self, key: K) -> None:
        """Evict a key from cache, calling callback if registered."""
        if key in self._cache:
            if self._on_evict:
                value, _ = self._cache[key]
                try:
                    self._on_evict(key, value)
                except Exception as e:
                    logger.warning(f"Eviction callback failed for {key}: {e}")
            del self._cache[key]

    @property
    def stats(self) -> dict[str, Any]:
        """Get cache statistics."""
        with self._lock:
            total = self._hits + self._misses
            hit_rate = self._hits / total if total > 0 else 0.0
            return {
                "size": len(self._cache),
                "max_size": self._max_size,
                "hits": self._hits,
                "misses": self._misses,
                "hit_rate": hit_rate,
                "ttl_seconds": self._ttl,
            }

    def __len__(self) -> int:
        with self._lock:
            return len(self._cache)

    def __contains__(self, key: K) -> bool:
        with self._lock:
            if key not in self._cache:
                return False
            _, timestamp = self._cache[key]
            return time.time() - timestamp <= self._ttl


# ============================================================================
# Content-Addressable Cache for Segment Predictions
# ============================================================================


def compute_file_hash(path: str | Path, algorithm: str = "sha256") -> str:
    """
    Compute hash of file contents for cache keying.

    Uses chunked reading for memory efficiency with large files.

    Args:
        path: Path to file
        algorithm: Hash algorithm (default: sha256)

    Returns:
        Hexadecimal hash string
    """
    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f"File not found: {path}")

    hasher = hashlib.new(algorithm)
    chunk_size = 8192  # 8KB chunks

    with open(path, "rb") as f:
        while chunk := f.read(chunk_size):
            hasher.update(chunk)

    return hasher.hexdigest()


def compute_quick_hash(path: str | Path) -> str:
    """
    Compute quick hash using file metadata (size + mtime).

    Much faster than content hash, suitable for most caching needs.

    Args:
        path: Path to file

    Returns:
        Hash string based on path, size, and modification time
    """
    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f"File not found: {path}")

    stat = path.stat()
    content = f"{path.resolve()}:{stat.st_size}:{stat.st_mtime}"
    return hashlib.sha256(content.encode()).hexdigest()[:16]


@dataclass
class SegmentPredictionCache:
    """
    Content-addressable cache for segment predictions.

    Caches predictions by feature file hash to avoid recomputation
    when the same features are used multiple times.

    Example:
        >>> cache = SegmentPredictionCache()
        >>> cache_key = cache.make_key("/path/to/features.h5", model_version="v1.0")
        >>> if cached := cache.get(cache_key):
        ...     return cached
        >>> # Compute predictions...
        >>> cache.set(cache_key, predictions)
    """

    max_size: int = 200
    ttl_seconds: float = 3600  # 1 hour
    _cache: LRUCache[Any, Any] = field(init=False, repr=False)

    def __post_init__(self) -> None:
        self._cache = LRUCache(max_size=self.max_size, ttl_seconds=self.ttl_seconds)

    def make_key(
        self,
        feature_path: str | Path,
        model_version: str = "default",
        params: dict[str, Any] | None = None,
    ) -> str:
        """
        Create cache key from feature file and model parameters.

        Args:
            feature_path: Path to feature file (uses quick hash)
            model_version: Model version identifier
            params: Additional parameters affecting prediction

        Returns:
            Cache key string
        """
        try:
            file_hash = compute_quick_hash(feature_path)
        except FileNotFoundError:
            # File doesn't exist, use path as key component
            file_hash = hashlib.sha256(str(feature_path).encode()).hexdigest()[:16]

        key_parts = [file_hash, model_version]
        if params:
            # Sort params for consistent hashing
            param_str = str(sorted(params.items()))
            param_hash = hashlib.sha256(param_str.encode()).hexdigest()[:8]
            key_parts.append(param_hash)

        return ":".join(key_parts)

    def get(self, key: str) -> list[dict[str, Any]] | None:
        """Get cached predictions by key."""
        return self._cache.get(key)

    def set(self, key: str, predictions: list[dict[str, Any]]) -> None:
        """Cache predictions with key."""
        self._cache.set(key, predictions)

    def invalidate(self, feature_path: str | Path) -> None:
        """Invalidate all cache entries for a feature file."""
        try:
            file_hash = compute_quick_hash(feature_path)
        except FileNotFoundError:
            file_hash = hashlib.sha256(str(feature_path).encode()).hexdigest()[:16]

        # Note: This is O(n) - for production, consider a secondary index
        keys_to_delete = []
        with self._cache._lock:
            for key in self._cache._cache.keys():
                if key.startswith(file_hash):
                    keys_to_delete.append(key)

        for key in keys_to_delete:
            self._cache.delete(key)

    @property
    def stats(self) -> dict[str, Any]:
        """Get cache statistics."""
        return self._cache.stats


# ============================================================================
# Metric Computation Cache
# ============================================================================


@dataclass
class MetricCacheKey:
    """Key for metric cache based on file pair."""

    predictions_hash: str
    references_hash: str
    metric_name: str
    params_hash: str

    def __hash__(self):
        return hash(
            (self.predictions_hash, self.references_hash, self.metric_name, self.params_hash)
        )

    def __eq__(self, other):
        if not isinstance(other, MetricCacheKey):
            return False
        return (
            self.predictions_hash == other.predictions_hash
            and self.references_hash == other.references_hash
            and self.metric_name == other.metric_name
            and self.params_hash == other.params_hash
        )


@dataclass
class MetricCache:
    """
    Cache for computed metrics with automatic invalidation.

    Caches evaluation results by prediction/reference file pair hash.

    Example:
        >>> cache = MetricCache()
        >>> key = cache.make_key(
        ...     predictions_path="/path/to/preds.eaf",
        ...     references_path="/path/to/refs.eaf",
        ...     metric_name="segmentation",
        ...     params={"tolerance_sec": 0.1}
        ... )
        >>> if cached := cache.get(key):
        ...     return cached
        >>> # Compute metrics...
        >>> cache.set(key, metrics)
    """

    max_size: int = 500
    ttl_seconds: float = 1800  # 30 minutes
    _cache: LRUCache[Any, Any] = field(init=False, repr=False)

    def __post_init__(self) -> None:
        self._cache = LRUCache(max_size=self.max_size, ttl_seconds=self.ttl_seconds)

    def make_key(
        self,
        predictions_path: str | Path,
        references_path: str | Path,
        metric_name: str,
        params: dict[str, Any] | None = None,
    ) -> MetricCacheKey:
        """
        Create cache key from file paths and parameters.

        Args:
            predictions_path: Path to predictions file
            references_path: Path to references file
            metric_name: Name of metric being computed
            params: Parameters affecting metric computation

        Returns:
            MetricCacheKey for cache lookup
        """
        try:
            pred_hash = compute_quick_hash(predictions_path)
        except FileNotFoundError:
            pred_hash = hashlib.sha256(str(predictions_path).encode()).hexdigest()[:16]

        try:
            ref_hash = compute_quick_hash(references_path)
        except FileNotFoundError:
            ref_hash = hashlib.sha256(str(references_path).encode()).hexdigest()[:16]

        params_hash = ""
        if params:
            param_str = str(sorted(params.items()))
            params_hash = hashlib.sha256(param_str.encode()).hexdigest()[:8]

        return MetricCacheKey(
            predictions_hash=pred_hash,
            references_hash=ref_hash,
            metric_name=metric_name,
            params_hash=params_hash,
        )

    def get(self, key: MetricCacheKey) -> dict[str, float] | None:
        """Get cached metrics by key."""
        return self._cache.get(key)

    def set(self, key: MetricCacheKey, metrics: dict[str, float]) -> None:
        """Cache metrics with key."""
        self._cache.set(key, metrics)

    def invalidate_by_file(self, path: str | Path) -> int:
        """
        Invalidate all cache entries involving a file.

        Returns number of entries invalidated.
        """
        try:
            file_hash = compute_quick_hash(path)
        except FileNotFoundError:
            file_hash = hashlib.sha256(str(path).encode()).hexdigest()[:16]

        count = 0
        keys_to_delete = []
        with self._cache._lock:
            for key in self._cache._cache.keys():
                if isinstance(key, MetricCacheKey):
                    if key.predictions_hash == file_hash or key.references_hash == file_hash:
                        keys_to_delete.append(key)

        for key in keys_to_delete:
            if self._cache.delete(key):
                count += 1

        return count

    @property
    def stats(self) -> dict[str, Any]:
        """Get cache statistics."""
        return self._cache.stats


# ============================================================================
# Batch Operations
# ============================================================================


@dataclass
class BatchUpdateResult:
    """Result of a batch segment update operation."""

    total: int
    successful: int
    failed: int
    errors: list[dict[str, str | int]] = field(default_factory=list)
    duration_ms: float = 0.0


def batch_update_segments(
    segments: list[dict[str, Any]],
    output_path: str | Path,
    video_path: str | Path | None = None,
    chunk_size: int = 100,
    progress_callback: Callable[[int, int], None] | None = None,
) -> BatchUpdateResult:
    """
    Efficiently update multiple segments in a single transaction.

    Optimized for updating 100+ segments with minimal I/O overhead.

    Args:
        segments: List of segment dictionaries with start, end, label, etc.
        output_path: Path to output ELAN file
        video_path: Optional video path for ELAN file
        chunk_size: Number of segments to process per chunk
        progress_callback: Optional callback(processed, total) for progress

    Returns:
        BatchUpdateResult with success/failure counts

    Example:
        >>> segments = [
        ...     {"start": 0.0, "end": 1.0, "label": "HELLO", "tier": "gloss"},
        ...     {"start": 1.5, "end": 2.5, "label": "WORLD", "tier": "gloss"},
        ... ]
        >>> result = batch_update_segments(segments, "/output/annotations.eaf")
        >>> print(f"Updated {result.successful}/{result.total} segments")
    """
    from sltk.data.core import Segment, SegmentList

    start_time = time.time()
    result = BatchUpdateResult(total=len(segments), successful=0, failed=0)

    if not segments:
        return result

    # Convert dictionaries to Segment objects in chunks
    segment_objects: list[Segment] = []
    for i in range(0, len(segments), chunk_size):
        chunk = segments[i : i + chunk_size]

        for j, seg_dict in enumerate(chunk):
            try:
                segment = Segment(
                    start=float(seg_dict.get("start", 0.0)),
                    end=float(seg_dict.get("end", 0.0)),
                    label=seg_dict.get("label"),
                    confidence=seg_dict.get("confidence"),
                    tier=seg_dict.get("tier", "default"),
                    metadata=seg_dict.get("metadata", {}),
                )
                segment_objects.append(segment)
                result.successful += 1
            except (ValueError, TypeError, KeyError) as e:
                result.failed += 1
                result.errors.append(
                    {"index": i + j, "error": str(e), "segment": str(seg_dict)[:100]}
                )

        if progress_callback:
            progress_callback(min(i + chunk_size, len(segments)), len(segments))

    # Write all segments at once
    if segment_objects:
        try:
            segment_list = SegmentList(segment_objects)
            segment_list.to_elan(output_path, video_path=video_path)
        except Exception as e:
            logger.error(f"Failed to write segments to {output_path}: {e}")
            # Mark all as failed
            result.failed = result.total
            result.successful = 0
            result.errors.append({"index": -1, "error": f"Write failed: {e}", "segment": ""})

    result.duration_ms = (time.time() - start_time) * 1000
    return result


@dataclass
class BatchMetricResult:
    """Result of a batch metric computation."""

    total_pairs: int
    successful: int
    failed: int
    results: list[dict[str, Any]] = field(default_factory=list)
    errors: list[dict[str, str]] = field(default_factory=list)
    duration_ms: float = 0.0
    cache_hits: int = 0


def batch_compute_metrics(
    segment_pairs: list[tuple[str, str]],
    metric_type: str = "segmentation",
    tolerance_sec: float = 0.1,
    cache: MetricCache | None = None,
    max_workers: int = 4,
    progress_callback: Callable[[int, int], None] | None = None,
) -> BatchMetricResult:
    """
    Compute metrics for multiple segment file pairs efficiently.

    Uses caching and parallel processing for optimal performance.

    Args:
        segment_pairs: List of (predictions_path, references_path) tuples
        metric_type: Type of metric to compute ("segmentation" or "translation")
        tolerance_sec: Tolerance for boundary matching (segmentation only)
        cache: Optional MetricCache for result caching
        max_workers: Maximum parallel workers
        progress_callback: Optional callback(processed, total) for progress

    Returns:
        BatchMetricResult with computed metrics

    Example:
        >>> pairs = [
        ...     ("/preds/file1.eaf", "/refs/file1.eaf"),
        ...     ("/preds/file2.eaf", "/refs/file2.eaf"),
        ... ]
        >>> result = batch_compute_metrics(pairs, metric_type="segmentation")
        >>> for r in result.results:
        ...     print(f"{r['predictions']}: F1={r['metrics']['boundary_f1']:.3f}")
    """
    import concurrent.futures

    from sltk.io.elan import read_eaf
    from sltk.metrics.segmentation import segmentation_metrics

    start_time = time.time()
    result = BatchMetricResult(total_pairs=len(segment_pairs), successful=0, failed=0)

    if not segment_pairs:
        return result

    def compute_single(pair: tuple[str, str]) -> dict[str, Any]:
        """Compute metrics for a single pair."""
        pred_path, ref_path = pair
        params = {"tolerance_sec": tolerance_sec}

        # Check cache first
        if cache:
            cache_key = cache.make_key(pred_path, ref_path, metric_type, params)
            if cached := cache.get(cache_key):
                return {
                    "predictions": pred_path,
                    "references": ref_path,
                    "metrics": cached,
                    "cached": True,
                }

        # Compute metrics
        try:
            predictions = read_eaf(pred_path)
            references = read_eaf(ref_path)

            if metric_type == "segmentation":
                metrics = segmentation_metrics(predictions, references, tolerance_sec=tolerance_sec)
            else:
                raise ValueError(f"Unknown metric type: {metric_type}")

            # Cache result
            if cache:
                cache.set(cache_key, metrics)

            return {
                "predictions": pred_path,
                "references": ref_path,
                "metrics": metrics,
                "cached": False,
            }
        except Exception as e:
            return {
                "predictions": pred_path,
                "references": ref_path,
                "error": str(e),
                "cached": False,
            }

    # Process pairs with thread pool
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {executor.submit(compute_single, pair): pair for pair in segment_pairs}

        for i, future in enumerate(concurrent.futures.as_completed(futures)):
            try:
                single_result = future.result()

                if "error" in single_result:
                    result.failed += 1
                    result.errors.append(
                        {
                            "predictions": single_result["predictions"],
                            "references": single_result["references"],
                            "error": single_result["error"],
                        }
                    )
                else:
                    result.successful += 1
                    result.results.append(single_result)
                    if single_result.get("cached"):
                        result.cache_hits += 1

            except Exception as e:
                result.failed += 1
                pair = futures[future]
                result.errors.append(
                    {"predictions": pair[0], "references": pair[1], "error": str(e)}
                )

            if progress_callback:
                progress_callback(i + 1, len(segment_pairs))

    result.duration_ms = (time.time() - start_time) * 1000
    return result


# ============================================================================
# Lazy Loading Utilities
# ============================================================================


@dataclass
class LazySegmentLoader:
    """
    Lazy loader for large segment collections.

    Loads segment details on demand to minimize memory usage.

    Example:
        >>> loader = LazySegmentLoader(segment_paths)
        >>> for batch in loader.iter_batches(batch_size=50):
        ...     process_segments(batch)
    """

    paths: list[str | Path]
    _loaded: dict[int, Any] = field(default_factory=dict, init=False)

    def __len__(self) -> int:
        return len(self.paths)

    def __getitem__(self, idx: int):
        """Load segment on demand."""
        if idx not in self._loaded:
            self._loaded[idx] = self._load_single(idx)
        return self._loaded[idx]

    def _load_single(self, idx: int):
        """Load a single segment file."""
        from sltk.io.elan import read_eaf

        path = self.paths[idx]
        return read_eaf(path)

    def iter_batches(
        self,
        batch_size: int = 50,
        preload: bool = True,
    ):
        """
        Iterate over segments in batches.

        Args:
            batch_size: Number of segments per batch
            preload: If True, preload next batch while processing current

        Yields:
            List of (index, segment_list) tuples
        """
        import concurrent.futures

        if not preload:
            # Simple sequential loading
            for i in range(0, len(self.paths), batch_size):
                batch = []
                for j in range(i, min(i + batch_size, len(self.paths))):
                    batch.append((j, self[j]))
                yield batch
                # Clear loaded to save memory
                self._loaded.clear()
            return

        # Preloading with thread pool
        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
            next_future = None
            next_start = 0

            for i in range(0, len(self.paths), batch_size):
                # Wait for preloaded batch if available
                if next_future and next_start == i:
                    batch = next_future.result()
                else:
                    batch = [
                        (j, self._load_single(j))
                        for j in range(i, min(i + batch_size, len(self.paths)))
                    ]

                # Start preloading next batch
                next_start = i + batch_size
                if next_start < len(self.paths):
                    next_future = executor.submit(
                        lambda start: [
                            (j, self._load_single(j))
                            for j in range(start, min(start + batch_size, len(self.paths)))
                        ],
                        next_start,
                    )
                else:
                    next_future = None

                yield batch
                self._loaded.clear()

    def unload(self, idx: int | None = None) -> None:
        """
        Unload segments to free memory.

        Args:
            idx: Specific index to unload, or None to unload all
        """
        if idx is None:
            self._loaded.clear()
        elif idx in self._loaded:
            del self._loaded[idx]


# ============================================================================
# Global Cache Instances
# ============================================================================

# Global cache instances - initialized lazily
_segment_prediction_cache: SegmentPredictionCache | None = None
_metric_cache: MetricCache | None = None
_cache_lock = threading.Lock()


def get_segment_prediction_cache() -> SegmentPredictionCache:
    """Get or create the global segment prediction cache."""
    global _segment_prediction_cache
    with _cache_lock:
        if _segment_prediction_cache is None:
            _segment_prediction_cache = SegmentPredictionCache(max_size=200, ttl_seconds=3600)
        return _segment_prediction_cache


def get_metric_cache() -> MetricCache:
    """Get or create the global metric cache."""
    global _metric_cache
    with _cache_lock:
        if _metric_cache is None:
            _metric_cache = MetricCache(max_size=500, ttl_seconds=1800)
        return _metric_cache


def clear_all_caches() -> dict[str, int]:
    """
    Clear all global caches and return sizes before clearing.

    Returns:
        Dictionary with cache sizes before clearing
    """
    global _segment_prediction_cache, _metric_cache
    sizes = {}

    with _cache_lock:
        if _segment_prediction_cache is not None:
            sizes["segment_predictions"] = len(_segment_prediction_cache._cache)
            _segment_prediction_cache._cache.clear()

        if _metric_cache is not None:
            sizes["metrics"] = len(_metric_cache._cache)
            _metric_cache._cache.clear()

    return sizes


def get_cache_stats() -> dict[str, Any]:
    """Get statistics for all global caches."""
    stats = {}

    with _cache_lock:
        if _segment_prediction_cache:
            stats["segment_predictions"] = _segment_prediction_cache.stats

        if _metric_cache:
            stats["metrics"] = _metric_cache.stats

    return stats
