src/dashboard/heatmap_source.py

"""Lazy dashboard heatmaps and safe snapshot selection.

This module provides memory-efficient heatmap generation from B5D snapshots
without loading the entire network into RAM. It uses memory-mapped I/O
and lazy projection for efficient visualization of structural data.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

from src.storage.b5d import B5DReader
from src.storage.lazy_view import B5DLazyProjector, StorageHeatmapKind

from .models import JSONValue

# ============================================================================
# Custom Exceptions
# ============================================================================


class HeatmapError(Exception):
    """Base exception for heatmap-related errors."""


class InvalidHeatmapKindError(HeatmapError):
    """Raised when an unknown heatmap kind is requested."""


class SnapshotNotFoundError(HeatmapError):
    """Raised when a requested snapshot file does not exist."""


class InvalidSnapshotError(HeatmapError):
    """Raised when a snapshot filename is invalid or unsafe."""


# ============================================================================
# Data Models
# ============================================================================


@dataclass(frozen=True, slots=True)
class HeatmapPayload:
    """JSON-ready heatmap projection.

    Attributes:
        kind: Type of heatmap (e.g., 'spike_activity', 'synaptic_weights').
        tick: The simulation tick at which the snapshot was taken.
        samples: Number of samples used in the projection.
        snapshot: Name of the snapshot file used.
        values: 2D array of heatmap values.
        source: Always "snapshot" to distinguish from live runtime data.
    """

    kind: str
    tick: int
    samples: int
    snapshot: str
    values: list[list[float]]
    source: str = "snapshot"

    def to_json(self) -> dict[str, JSONValue]:
        """Return the projection in JSON-compatible form."""
        return {
            "kind": self.kind,
            "tick": self.tick,
            "samples": self.samples,
            "snapshot": self.snapshot,
            "values": [list(row) for row in self.values],
            "source": self.source,
        }


@dataclass(frozen=True, slots=True)
class SnapshotEntry:
    """One selectable snapshot beside the configured default snapshot.

    Attributes:
        name: Filename of the snapshot.
        size_bytes: File size in bytes.
    """

    name: str
    size_bytes: int

    def __post_init__(self) -> None:
        """Validate snapshot entry data."""
        if not self.name or not self.name.endswith(".b5d"):
            raise ValueError(f"Invalid snapshot name: {self.name}")
        if self.size_bytes < 0:
            raise ValueError(f"Invalid file size: {self.size_bytes}")

    def to_json(self) -> dict[str, JSONValue]:
        """Return JSON-ready snapshot metadata."""
        return {"name": self.name, "size_bytes": self.size_bytes}


# ============================================================================
# Main Heatmap Source
# ============================================================================


class SnapshotHeatmapSource:
    """Build heatmaps without loading the whole network into RAM.

    This class provides memory-efficient heatmap generation from B5D snapshots
    using memory-mapped I/O and lazy projection. It supports listing available
    snapshots and building heatmaps from specific snapshots.

    The heatmap source is designed to be used with the dashboard for
    visualization of structural and activity patterns.

    Example:
        >>> source = SnapshotHeatmapSource(Path("snapshots/latest.b5d"))
        >>> snapshots = source.list_snapshots()
        >>> heatmap = source.build("spike_activity", snapshot_name="snapshot_1000.b5d")
    """

    def __init__(self, snapshot_path: Path, activity_tau_ticks: float = 50.0) -> None:
        """Initialize the heatmap source.

        Args:
            snapshot_path: Path to the default snapshot file.
            activity_tau_ticks: Time constant for activity decay in ticks.
                Higher values smooth activity over longer windows.
        """
        self.snapshot_path = snapshot_path.resolve()
        self.snapshot_root = self.snapshot_path.parent
        self.activity_tau_ticks = activity_tau_ticks

        if not self.snapshot_path.is_file():
            raise FileNotFoundError(f"Default snapshot not found: {self.snapshot_path}")

    def list_snapshots(self) -> tuple[SnapshotEntry, ...]:
        """List all available B5D snapshots in the snapshot directory.

        Returns:
            Tuple of SnapshotEntry objects, sorted by name (typically timestamp).
            Returns empty tuple if directory does not exist.

        Note:
            Only files with .b5d extension are included. Hidden files and
            directories are excluded.
        """
        if not self.snapshot_root.is_dir():
            return ()

        entries: list[SnapshotEntry] = []

        for path in sorted(self.snapshot_root.glob("*.b5d")):
            if not path.is_file():
                continue
            if path.name.startswith("."):
                continue

            try:
                entries.append(SnapshotEntry(path.name, path.stat().st_size))
            except (OSError, ValueError):
                # Skip files that cannot be read
                continue

        return tuple(entries)

    def build(self, kind_name: str, snapshot_name: str | None = None) -> HeatmapPayload:
        """Build a heatmap from a snapshot.

        Args:
            kind_name: Type of heatmap to build (e.g., 'spike_activity').
                Must be a valid StorageHeatmapKind value.
            snapshot_name: Optional specific snapshot filename. If None or
                empty, uses the default snapshot.

        Returns:
            HeatmapPayload containing the projection data.

        Raises:
            InvalidHeatmapKindError: If kind_name is not a valid heatmap type.
            SnapshotNotFoundError: If the requested snapshot file does not exist.
            InvalidSnapshotError: If the snapshot filename is invalid or unsafe.
            HeatmapError: For other heatmap-related errors.
        """
        # Validate heatmap kind
        try:
            kind = StorageHeatmapKind(kind_name)
        except ValueError as e:
            raise InvalidHeatmapKindError(
                f"Unknown heatmap kind: {kind_name}. "
                f"Valid types: {[k.value for k in StorageHeatmapKind]}"
            ) from e

        # Resolve snapshot path
        try:
            snapshot = self._resolve_snapshot(snapshot_name)
        except ValueError as e:
            raise InvalidSnapshotError(str(e)) from e
        except FileNotFoundError as e:
            raise SnapshotNotFoundError(str(e)) from e

        # Build the heatmap using lazy projection
        try:
            with B5DReader(snapshot) as reader:
                projector = B5DLazyProjector(reader, self.activity_tau_ticks)
                result = projector.build(kind)

            return HeatmapPayload(
                kind=result.kind.value,
                tick=result.snapshot_tick,
                samples=result.samples,
                snapshot=snapshot.name,
                values=result.values.tolist(),
            )

        except Exception as e:
            raise HeatmapError(f"Failed to build heatmap: {e}") from e

    def _resolve_snapshot(self, snapshot_name: str | None) -> Path:
        """Resolve a snapshot path with security checks.

        Args:
            snapshot_name: Name of the snapshot file, or None for default.

        Returns:
            Resolved Path object.

        Raises:
            ValueError: If the snapshot name is invalid or unsafe.
            FileNotFoundError: If the snapshot file does not exist.
        """
        # Use default if no snapshot name provided
        if snapshot_name is None or not snapshot_name.strip():
            return self.snapshot_path

        # Basic security: prevent path traversal
        if ".." in snapshot_name:
            raise ValueError(f"Invalid snapshot filename: {snapshot_name}")

        # Ensure the filename is clean
        if Path(snapshot_name).name != snapshot_name:
            raise ValueError(f"Invalid snapshot filename: {snapshot_name}")

        # Ensure it's a .b5d file
        if not snapshot_name.endswith(".b5d"):
            raise ValueError(f"Invalid snapshot extension: {snapshot_name}")

        # Resolve and ensure it's within the snapshot root
        candidate = (self.snapshot_root / snapshot_name).resolve()

        if candidate.parent != self.snapshot_root:
            raise ValueError(f"Snapshot not in allowed directory: {snapshot_name}")

        if not candidate.is_file():
            raise FileNotFoundError(f"Snapshot not found: {snapshot_name}")

        return candidate

    def has_snapshot(self, snapshot_name: str) -> bool:
        """Check if a specific snapshot exists.

        Args:
            snapshot_name: Name of the snapshot file.

        Returns:
            True if the snapshot exists and is accessible.
        """
        try:
            candidate = self._resolve_snapshot(snapshot_name)
            return candidate.is_file()
        except (ValueError, FileNotFoundError):
            return False

    def get_snapshot_info(self, snapshot_name: str) -> SnapshotEntry:
        """Get metadata for a specific snapshot.

        Args:
            snapshot_name: Name of the snapshot file.

        Returns:
            SnapshotEntry with file metadata.

        Raises:
            ValueError: If the snapshot name is invalid.
            FileNotFoundError: If the snapshot file does not exist.
        """
        snapshot = self._resolve_snapshot(snapshot_name)
        return SnapshotEntry(snapshot.name, snapshot.stat().st_size)


# ============================================================================
# Demo Heatmap Source (Fallback when no .b5d snapshot exists)
# ============================================================================


class DemoHeatmapSource:
    """Fallback heatmap source that generates synthetic demo data.

    This source is used when no .b5d snapshot file is available. It
    generates realistic-looking synthetic heatmaps so the dashboard
    always has something to display, even without a running simulation.
    """

    def __init__(self) -> None:
        self._tick = 0
        self._rng_seed = 42

    def list_snapshots(self) -> tuple[SnapshotEntry, ...]:
        """Return an empty list (no real snapshots available)."""
        return ()

    def build(
        self,
        kind_name: str,
        _snapshot_name: str | None = None,
    ) -> HeatmapPayload:
        """Build a synthetic demo heatmap.

        Args:
            kind_name: Type of heatmap (ignored for demo).
            snapshot_name: Optional snapshot name (ignored for demo).

        Returns:
            A HeatmapPayload with synthetic data.
        """
        self._tick += 1
        import random

        rng = random.Random(self._rng_seed + self._tick)

        # Generate a 40x20 grid of synthetic values with some structure
        rows, cols = 40, 20
        values: list[list[float]] = []
        for y in range(cols):
            row: list[float] = []
            for x in range(rows):
                # Create some spatial structure: clusters of activity
                cx, cy = rows // 2, cols // 2
                dist = ((x - cx) / cx) ** 2 + ((y - cy) / cy) ** 2
                cluster = max(0, 1.0 - dist * 0.5)
                noise = rng.gauss(0, 0.15)
                val = max(0.0, min(1.0, cluster + noise))
                row.append(val)
            values.append(row)

        return HeatmapPayload(
            kind=kind_name or "activity",
            tick=self._tick * 100,
            samples=rows * cols,
            snapshot="demo (no .b5d snapshot)",
            values=values,
        )

    def has_snapshot(self, _snapshot_name: str) -> bool:
        """Demo source has no real snapshots."""
        return False

    def get_snapshot_info(self, snapshot_name: str) -> SnapshotEntry:
        """Demo source has no real snapshots."""
        raise FileNotFoundError(f"No snapshot available: {snapshot_name}")


# ============================================================================
# Factory Function
# ============================================================================


def create_heatmap_source(
    snapshot_path: str | Path,
    activity_tau_ticks: float = 50.0,
) -> SnapshotHeatmapSource:
    """Factory function for creating a SnapshotHeatmapSource.

    Args:
        snapshot_path: Path to the default snapshot file.
        activity_tau_ticks: Time constant for activity decay.

    Returns:
        Configured SnapshotHeatmapSource instance.

    Example:
        >>> source = create_heatmap_source("snapshots/latest.b5d")
        >>> source.list_snapshots()
    """
    if isinstance(snapshot_path, str):
        snapshot_path = Path(snapshot_path)
    return SnapshotHeatmapSource(snapshot_path, activity_tau_ticks)