#!/usr/bin/env python3
"""
bar_lru.py - BAR1 LRU Manager for Zero-Copy GPU Memory

Manages a budget of GPU-registered host memory with LRU eviction.
When the budget is exceeded, least-recently-used spans are evicted.

Key concepts:
  - BAR1: Base Address Register 1 - GPU's window into host memory
  - Registered spans: Host memory registered with cudaHostRegister
  - Budget: Maximum total bytes to keep registered
  - LRU eviction: Remove oldest spans when budget exceeded

Work Order: WO-ZC-SHIM

Usage:
    from helix_cdc.zc.bar_lru import BARLRU

    lru = BARLRU(budget_bytes=512 * 1024 * 1024)  # 512 MB budget

    # Acquire a span (may evict old ones)
    span, created = lru.acquire("/path/to/shard.cdna2.hxz", offset=0, nbytes=1024*1024)

    # Access data via span.host_ptr or span.dev_alias
    # ...

    # Stats
    print(lru.stats())
"""

from __future__ import annotations

import threading
import time
from collections import OrderedDict
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from .zc_bind import HostSpan, map_register, unregister_unmap


@dataclass
class CachedSpan:
    """Cached HostSpan with metadata."""
    span: HostSpan
    path: str
    offset: int
    nbytes: int
    created_at: float = field(default_factory=time.time)
    last_access: float = field(default_factory=time.time)
    access_count: int = 0

    def touch(self) -> None:
        """Update access time and count."""
        self.last_access = time.time()
        self.access_count += 1

    @property
    def key(self) -> str:
        """Unique key for this span."""
        return f"{self.path}:{self.offset}:{self.nbytes}"


class BARLRU:
    """
    LRU cache for GPU-registered memory spans.

    Manages a budget of registered memory. When budget is exceeded,
    evicts least-recently-used spans to make room.

    Thread-safe.
    """

    def __init__(
        self,
        budget_bytes: int = 512 * 1024 * 1024,  # 512 MB default
        min_evict_age_seconds: float = 1.0,     # Don't evict spans younger than this
    ):
        """
        Initialize BAR1 LRU manager.

        Args:
            budget_bytes: Maximum total registered bytes
            min_evict_age_seconds: Minimum age before span can be evicted
        """
        self.budget_bytes = budget_bytes
        self.min_evict_age_seconds = min_evict_age_seconds

        self._cache: OrderedDict[str, CachedSpan] = OrderedDict()
        self._current_bytes: int = 0
        self._lock = threading.RLock()

        # Statistics
        self._hits: int = 0
        self._misses: int = 0
        self._evictions: int = 0
        self._total_registered: int = 0
        self._total_unregistered: int = 0

    def acquire(
        self,
        path: str,
        offset: int,
        nbytes: int,
        onfault: bool = True,
    ) -> Tuple[HostSpan, bool]:
        """
        Acquire a registered memory span.

        If the span is already cached, return it and update LRU.
        Otherwise, register new span (evicting old ones if needed).

        Args:
            path: File path
            offset: Byte offset in file
            nbytes: Number of bytes to map
            onfault: Eagerly fault pages

        Returns:
            (HostSpan, created_bool) where created_bool is True if newly created

        Raises:
            OSError: On registration failure
        """
        key = f"{path}:{offset}:{nbytes}"

        with self._lock:
            # Check cache
            if key in self._cache:
                cached = self._cache[key]
                cached.touch()

                # Move to end (most recently used)
                self._cache.move_to_end(key)

                self._hits += 1
                return cached.span, False

            # Cache miss - need to register
            self._misses += 1

            # Evict if needed to make room
            self._evict_until_fits(nbytes)

            # Register new span
            span = map_register(path, offset, nbytes, onfault)

            # Add to cache
            cached = CachedSpan(
                span=span,
                path=path,
                offset=offset,
                nbytes=nbytes,
            )
            self._cache[key] = cached
            self._current_bytes += nbytes
            self._total_registered += 1

            return span, True

    def _evict_until_fits(self, needed_bytes: int) -> int:
        """
        Evict LRU spans until there's room for needed_bytes.

        Returns number of spans evicted.
        """
        evicted = 0
        now = time.time()

        while self._current_bytes + needed_bytes > self.budget_bytes:
            if not self._cache:
                break  # Nothing left to evict

            # Get oldest (first in OrderedDict)
            oldest_key = next(iter(self._cache))
            oldest = self._cache[oldest_key]

            # Check minimum age
            age = now - oldest.created_at
            if age < self.min_evict_age_seconds and self._current_bytes > 0:
                # Can't evict too-young spans, but need room
                # Force evict anyway if we're over budget
                if self._current_bytes <= self.budget_bytes:
                    break

            # Evict
            try:
                unregister_unmap(oldest.span)
            except OSError:
                pass  # Best effort cleanup

            del self._cache[oldest_key]
            self._current_bytes -= oldest.nbytes
            self._evictions += 1
            self._total_unregistered += 1
            evicted += 1

        return evicted

    def release(self, path: str, offset: int, nbytes: int) -> bool:
        """
        Explicitly release a span from the cache.

        Args:
            path: File path
            offset: Byte offset
            nbytes: Number of bytes

        Returns:
            True if span was found and released, False otherwise
        """
        key = f"{path}:{offset}:{nbytes}"

        with self._lock:
            if key not in self._cache:
                return False

            cached = self._cache[key]

            try:
                unregister_unmap(cached.span)
            except OSError:
                pass

            del self._cache[key]
            self._current_bytes -= cached.nbytes
            self._total_unregistered += 1

            return True

    def clear(self) -> int:
        """
        Clear all cached spans.

        Returns:
            Number of spans cleared
        """
        with self._lock:
            count = len(self._cache)

            for cached in self._cache.values():
                try:
                    unregister_unmap(cached.span)
                except OSError:
                    pass

            self._cache.clear()
            self._current_bytes = 0
            self._total_unregistered += count

            return count

    def contains(self, path: str, offset: int, nbytes: int) -> bool:
        """Check if span is in cache."""
        key = f"{path}:{offset}:{nbytes}"
        with self._lock:
            return key in self._cache

    def stats(self) -> Dict[str, Any]:
        """
        Get cache statistics.

        Returns:
            Dictionary with stats
        """
        with self._lock:
            entries = []
            for cached in self._cache.values():
                entries.append({
                    "path": cached.path,
                    "offset": cached.offset,
                    "nbytes": cached.nbytes,
                    "access_count": cached.access_count,
                    "age_seconds": time.time() - cached.created_at,
                })

            return {
                "budget_bytes": self.budget_bytes,
                "current_bytes": self._current_bytes,
                "utilization": self._current_bytes / self.budget_bytes if self.budget_bytes > 0 else 0,
                "num_spans": len(self._cache),
                "hits": self._hits,
                "misses": self._misses,
                "hit_rate": self._hits / (self._hits + self._misses) if (self._hits + self._misses) > 0 else 0,
                "evictions": self._evictions,
                "total_registered": self._total_registered,
                "total_unregistered": self._total_unregistered,
                "entries": entries,
            }

    def to_receipt(self) -> Dict[str, Any]:
        """Generate a receipt for proof/audit."""
        stats = self.stats()
        return {
            "schema": "bar_lru_stats_v1",
            "timestamp_utc": datetime.now(timezone.utc).isoformat(),
            "budget_mb": self.budget_bytes / (1024 * 1024),
            "current_mb": self._current_bytes / (1024 * 1024),
            "utilization_percent": stats["utilization"] * 100,
            "num_spans": stats["num_spans"],
            "cache_stats": {
                "hits": self._hits,
                "misses": self._misses,
                "hit_rate_percent": stats["hit_rate"] * 100,
                "evictions": self._evictions,
            },
            "lifetime_stats": {
                "total_registered": self._total_registered,
                "total_unregistered": self._total_unregistered,
            },
        }

    def __len__(self) -> int:
        """Return number of cached spans."""
        with self._lock:
            return len(self._cache)

    def __enter__(self) -> "BARLRU":
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        """Context manager exit - clear all spans."""
        self.clear()


# Global singleton for convenience
_global_lru: Optional[BARLRU] = None


def get_global_lru(budget_bytes: int = 512 * 1024 * 1024) -> BARLRU:
    """
    Get or create global BAR1 LRU instance.

    Args:
        budget_bytes: Budget for new instance (ignored if already exists)

    Returns:
        Global BARLRU instance
    """
    global _global_lru
    if _global_lru is None:
        _global_lru = BARLRU(budget_bytes=budget_bytes)
    return _global_lru


def reset_global_lru() -> None:
    """Reset (clear and delete) global LRU instance."""
    global _global_lru
    if _global_lru is not None:
        _global_lru.clear()
        _global_lru = None


if __name__ == "__main__":
    import json

    print("bar_lru.py - BAR1 LRU Manager")
    print()

    # Demo
    lru = BARLRU(budget_bytes=1024 * 1024)  # 1 MB budget
    print(f"Budget: {lru.budget_bytes / (1024*1024):.1f} MB")
    print()

    print("Stats (empty):")
    print(json.dumps(lru.stats(), indent=2))
