backend/storage.py

"""
Zarr storage layer for efficient tensor serialization.

Stores instrumentation data to disk using Zarr with Blosc compression:
- Attention tensors: chunked by (layer, head) for fast slice access
- Residual norms, logits: standard chunking
- Metadata: JSON files

Storage structure:
/tmp/runs/{run_id}/
├── tensors/
│   ├── attention.zarr/
│   ├── residuals.zarr/
│   └── logits.zarr/
├── metadata.json
└── telemetry.jsonl
"""

import zarr
import numcodecs
import numpy as np
import torch
import json
import os
import shutil
from typing import Dict, Any, Optional, List
from pathlib import Path
from datetime import datetime
import logging

from .instrumentation import InstrumentationData, TokenMetadata, LayerMetadata

logger = logging.getLogger(__name__)


class ZarrStorage:
    """
    Manages Zarr storage for instrumentation data.

    Features:
    - Blosc compression (>3x compression ratio)
    - Chunking optimized for visualization access patterns
    - Lazy loading support
    - Export to zip bundles for study reproducibility
    """

    def __init__(self, run_id: str, base_dir: str = "/tmp/runs"):
        self.run_id = run_id
        self.base_dir = Path(base_dir)
        self.run_dir = self.base_dir / run_id
        self.tensor_dir = self.run_dir / "tensors"

        # Create directories
        self.tensor_dir.mkdir(parents=True, exist_ok=True)

        # Blosc compressor for efficient compression
        self.compressor = numcodecs.Blosc(
            cname='zstd',      # zstd algorithm (good compression + speed)
            clevel=5,          # Compression level (1-9, 5 is balanced)
            shuffle=numcodecs.Blosc.SHUFFLE  # Byte shuffle for better compression
        )

    def save_instrumentation_data(self, data: InstrumentationData) -> Dict[str, Any]:
        """
        Save complete instrumentation data to Zarr + JSON.

        Args:
            data: InstrumentationData from ModelInstrumentor

        Returns:
            Dictionary with file paths and sizes
        """
        logger.info(f"Saving instrumentation data for run {self.run_id}...")

        result = {
            'run_id': self.run_id,
            'paths': {},
            'sizes_mb': {}
        }

        # 1. Save attention tensors (largest data)
        if data.attention_tensors is not None:
            attn_path = self._save_attention_tensors(data.attention_tensors)
            result['paths']['attention'] = str(attn_path)
            result['sizes_mb']['attention'] = self._get_dir_size_mb(attn_path)

        # 2. Save metadata (JSON)
        metadata_path = self._save_metadata(data)
        result['paths']['metadata'] = str(metadata_path)
        result['sizes_mb']['metadata'] = self._get_file_size_mb(metadata_path)

        # 3. Save token data (JSON)
        tokens_path = self._save_token_data(data.tokens)
        result['paths']['tokens'] = str(tokens_path)
        result['sizes_mb']['tokens'] = self._get_file_size_mb(tokens_path)

        # 4. Save layer metadata (JSON)
        layer_meta_path = self._save_layer_metadata(data.layer_metadata)
        result['paths']['layer_metadata'] = str(layer_meta_path)
        result['sizes_mb']['layer_metadata'] = self._get_file_size_mb(layer_meta_path)

        # Summary
        total_size = sum(result['sizes_mb'].values())
        result['total_size_mb'] = total_size

        logger.info(f"✅ Saved {total_size:.2f} MB to {self.run_dir}")

        return result

    def _save_attention_tensors(self, attention_tensor: torch.Tensor) -> Path:
        """
        Save attention tensors with optimal chunking.

        Input shape: [num_tokens, num_layers, num_heads, seq_len, seq_len]
        Chunking: (1, 1, 1, seq_len, seq_len) - one chunk per layer/head

        This allows fast loading of individual head attention without
        loading the entire tensor.
        """
        path = self.tensor_dir / "attention.zarr"

        # Convert to numpy (Zarr doesn't support torch tensors directly)
        attn_np = attention_tensor.cpu().numpy()

        # Determine chunk shape
        num_tokens, num_layers, num_heads, seq_len, _ = attn_np.shape
        chunk_shape = (1, 1, 1, seq_len, seq_len)  # One chunk per layer/head

        # Save with compression
        z = zarr.open(
            str(path),
            mode='w',
            shape=attn_np.shape,
            chunks=chunk_shape,
            dtype=attn_np.dtype,
            compressor=self.compressor
        )
        z[:] = attn_np

        logger.info(f"   Attention: shape={attn_np.shape}, chunks={chunk_shape}")

        return path

    def _save_metadata(self, data: InstrumentationData) -> Path:
        """Save run metadata as JSON"""
        path = self.run_dir / "metadata.json"

        metadata = {
            'run_id': data.run_id,
            'seed': data.seed,
            'model_name': data.model_name,
            'timestamp': data.timestamp,
            'timestamp_iso': datetime.fromtimestamp(data.timestamp).isoformat(),
            'prompt': data.prompt,
            'max_tokens': data.max_tokens,
            'temperature': data.temperature,
            'top_k': data.top_k,
            'top_p': data.top_p,
            'total_time_ms': data.total_time_ms,
            'num_layers': data.num_layers,
            'num_heads': data.num_heads,
            'seq_length': data.seq_length,
            'num_generated_tokens': len(data.tokens)
        }

        with open(path, 'w') as f:
            json.dump(metadata, f, indent=2)

        return path

    def _save_token_data(self, tokens: List[TokenMetadata]) -> Path:
        """Save token metadata as JSON"""
        path = self.run_dir / "tokens.json"

        tokens_data = [
            {
                'token_id': t.token_id,
                'text': t.text,
                'position': t.position,
                'logprob': t.logprob,
                'entropy': t.entropy,
                'top_k_tokens': t.top_k_tokens,
                'byte_length': t.byte_length,
                'timestamp_ms': t.timestamp_ms
            }
            for t in tokens
        ]

        with open(path, 'w') as f:
            json.dump(tokens_data, f, indent=2)

        return path

    def _save_layer_metadata(self, layer_metadata: List[List[LayerMetadata]]) -> Path:
        """Save layer-level metadata as JSON"""
        path = self.run_dir / "layer_metadata.json"

        # Convert to serializable format
        layer_data = [
            [
                {
                    'layer_idx': lm.layer_idx,
                    'residual_norm': lm.residual_norm,
                    'time_ms': lm.time_ms,
                    'attention_output_norm': lm.attention_output_norm,
                    'ffn_output_norm': lm.ffn_output_norm
                }
                for lm in token_layers
            ]
            for token_layers in layer_metadata
        ]

        with open(path, 'w') as f:
            json.dump(layer_data, f, indent=2)

        return path

    def load_attention_slice(self, layer_idx: int, head_idx: int, token_idx: int = 0) -> np.ndarray:
        """
        Load a single attention head's matrix for a specific token.

        Args:
            layer_idx: Layer index (0-31 for Code Llama)
            head_idx: Head index (0-31 for Code Llama)
            token_idx: Token generation step (default 0 = first token)

        Returns:
            Attention matrix [seq_len, seq_len]
        """
        path = self.tensor_dir / "attention.zarr"

        if not path.exists():
            raise FileNotFoundError(f"Attention data not found at {path}")

        # Open in read mode
        z = zarr.open(str(path), mode='r')

        # Load specific slice
        # Shape: [num_tokens, num_layers, num_heads, seq_len, seq_len]
        attention_matrix = z[token_idx, layer_idx, head_idx, :, :]

        return attention_matrix

    def load_metadata(self) -> Dict[str, Any]:
        """Load run metadata"""
        path = self.run_dir / "metadata.json"
        with open(path, 'r') as f:
            return json.load(f)

    def load_tokens(self) -> List[Dict[str, Any]]:
        """Load token metadata"""
        path = self.run_dir / "tokens.json"
        with open(path, 'r') as f:
            return json.load(f)

    def export_bundle(self, output_path: Optional[Path] = None) -> Path:
        """
        Create a zip bundle of the entire run directory for export.

        Args:
            output_path: Optional custom output path (default: /tmp/run_{run_id}.zip)

        Returns:
            Path to created zip file
        """
        if output_path is None:
            output_path = self.base_dir / f"run_{self.run_id}.zip"

        logger.info(f"Creating export bundle: {output_path}")

        # Create zip archive
        shutil.make_archive(
            str(output_path.with_suffix('')),  # Remove .zip, make_archive adds it
            'zip',
            self.run_dir
        )

        bundle_size_mb = self._get_file_size_mb(output_path)
        logger.info(f"✅ Created bundle: {bundle_size_mb:.2f} MB")

        return output_path

    def cleanup(self):
        """Delete run directory and all contents"""
        if self.run_dir.exists():
            shutil.rmtree(self.run_dir)
            logger.info(f"Cleaned up run directory: {self.run_dir}")

    def _get_dir_size_mb(self, path: Path) -> float:
        """Get total size of directory in MB"""
        total_size = sum(
            f.stat().st_size for f in path.rglob('*') if f.is_file()
        )
        return total_size / (1024 * 1024)

    def _get_file_size_mb(self, path: Path) -> float:
        """Get file size in MB"""
        return path.stat().st_size / (1024 * 1024)


def generate_run_id() -> str:
    """
    Generate unique Run ID.

    Format: R{YYYY-MM-DD}-{HHMM}-{hash}
    Example: R2025-11-01-1430-a7f3
    """
    now = datetime.now()
    date_str = now.strftime("%Y-%m-%d")
    time_str = now.strftime("%H%M")

    # Short hash from timestamp microseconds
    hash_str = hex(now.microsecond)[-4:]

    return f"R{date_str}-{time_str}-{hash_str}"


def create_telemetry_log(run_id: str, base_dir: str = "/tmp/runs") -> Path:
    """
    Create telemetry JSONL file for logging events.

    Returns path to telemetry file.
    """
    run_dir = Path(base_dir) / run_id
    run_dir.mkdir(parents=True, exist_ok=True)

    telemetry_path = run_dir / "telemetry.jsonl"

    # Initialize with run.start event
    with open(telemetry_path, 'w') as f:
        f.write(json.dumps({
            'event': 'run.start',
            'run_id': run_id,
            'timestamp': datetime.now().timestamp()
        }) + '\n')

    return telemetry_path


def log_telemetry_event(run_id: str, event: str, data: Dict[str, Any],
                        base_dir: str = "/tmp/runs"):
    """
    Append telemetry event to JSONL log.

    Args:
        run_id: Run identifier
        event: Event name (e.g., 'token.emit', 'ablation.run')
        data: Event-specific data
        base_dir: Base directory for runs
    """
    telemetry_path = Path(base_dir) / run_id / "telemetry.jsonl"

    event_data = {
        'event': event,
        'timestamp': datetime.now().timestamp(),
        **data
    }

    with open(telemetry_path, 'a') as f:
        f.write(json.dumps(event_data) + '\n')


# Example usage
if __name__ == "__main__":
    print("Storage module loaded successfully")

    # Example: Create a mock run
    run_id = generate_run_id()
    print(f"Generated Run ID: {run_id}")

    storage = ZarrStorage(run_id)
    print(f"Storage directory: {storage.run_dir}")