#!/usr/bin/env python3
"""
gguf_fuse.py - Virtual model.gguf FUSE mount for Helix containers
=================================================================

Exposes Helix capsules as a virtual model.gguf file that can be read
directly by llama.cpp or other GGUF consumers.

Supported backends:
- hxv1: HXV1 vault with SVD-compressed layers
- npz: Directory of .npz encoded layers (from GGUF orchestrator)

Layout:
  /<mount>/
    model.gguf     # Virtual GGUF file (on-demand reconstruction)
    meta.json      # Container metadata
    layers/        # Individual layer files

Usage:
  python -m helix_cdc.fuse.gguf_fuse --hxv1 /path/to/model.vault --mount /mnt/gguf
  python -m helix_cdc.fuse.gguf_fuse --npz /path/to/encoded --mount /mnt/gguf
"""

import os
import sys
import json
import errno
import stat
import time
import hashlib
import argparse
from pathlib import Path
from typing import Dict, List, Optional, Any
import numpy as np

# Try to import fusepy
try:
    from fuse import FUSE, Operations, FuseOSError
    HAVE_FUSE = True
except ImportError:
    HAVE_FUSE = False
    # Stub for testing without fuse
    class Operations:
        pass
    class FuseOSError(Exception):
        def __init__(self, errno):
            self.errno = errno


class HXV1Backend:
    """Backend for HXV1 vault files."""

    def __init__(self, vault_path: Path):
        self.vault_path = vault_path
        self.K = 0
        self.D = 0
        self.basis = None
        self._load()

    def _load(self):
        """Load vault header and basis."""
        with open(self.vault_path, 'rb') as f:
            magic = f.read(4)
            if magic != b'HXV1':
                raise ValueError(f"Invalid vault magic: {magic}")

            import struct
            self.K = struct.unpack('<I', f.read(4))[0]
            self.D = struct.unpack('<I', f.read(4))[0]
            dtype_code = struct.unpack('<B', f.read(1))[0]
            f.read(3)  # padding

            dtype_map = {0: np.float32, 1: np.float64, 2: np.float16}
            dtype = dtype_map.get(dtype_code, np.float32)

            self.basis = np.frombuffer(f.read(self.K * self.D * dtype().itemsize), dtype=dtype)
            self.basis = self.basis.reshape(self.K, self.D)

    @property
    def num_layers(self) -> int:
        return self.K

    @property
    def layer_dim(self) -> int:
        return self.D

    def get_layer(self, idx: int) -> np.ndarray:
        """Get a single basis vector."""
        if idx < 0 or idx >= self.K:
            raise IndexError(f"Layer index {idx} out of range [0, {self.K})")
        return self.basis[idx]

    def reconstruct_tensor(self, coeffs: np.ndarray) -> np.ndarray:
        """Reconstruct tensor from coefficients."""
        return coeffs @ self.basis


class NPZBackend:
    """Backend for NPZ-encoded layers (from GGUF orchestrator)."""

    def __init__(self, npz_dir: Path):
        self.npz_dir = npz_dir
        self.manifest = None
        self.layers_dir = npz_dir / 'layers'
        self._load_manifest()

    def _load_manifest(self):
        """Load manifest.json."""
        manifest_path = self.npz_dir / 'manifest.json'
        if manifest_path.exists():
            self.manifest = json.loads(manifest_path.read_text())
        else:
            # Generate manifest from layer files
            self.manifest = {
                'schema': 'GGUF_ORCHESTRATOR_V1',
                'layers': []
            }
            if self.layers_dir.exists():
                for npz_file in sorted(self.layers_dir.glob('*.npz')):
                    self.manifest['layers'].append({
                        'name': npz_file.stem.replace('_', '.'),
                        'file': npz_file.name
                    })

    @property
    def num_layers(self) -> int:
        return len(self.manifest.get('layers', []))

    def get_layer_info(self, idx: int) -> Dict[str, Any]:
        """Get layer metadata."""
        layers = self.manifest.get('layers', [])
        if idx < 0 or idx >= len(layers):
            raise IndexError(f"Layer index {idx} out of range")
        return layers[idx]

    def reconstruct_layer(self, idx: int) -> np.ndarray:
        """Reconstruct a tensor from SVD components."""
        layer_info = self.get_layer_info(idx)
        npz_path = self.layers_dir / layer_info['file']

        data = np.load(npz_path)

        if 'U_k' in data and 'S_k' in data and 'Vt_k' in data:
            # SVD reconstruction: W = U @ diag(S) @ Vt
            U_k = data['U_k']
            S_k = data['S_k']
            Vt_k = data['Vt_k']
            return U_k @ np.diag(S_k) @ Vt_k
        elif 'tensor' in data:
            return data['tensor']
        else:
            raise ValueError(f"Unknown NPZ format in {npz_path}")


class HelixGGUFFuse(Operations):
    """
    FUSE filesystem exposing Helix containers as virtual model.gguf.

    This enables llama.cpp and other GGUF tools to read directly from
    compressed Helix representations without full decompression.
    """

    def __init__(self, backend_path: str, backend_type: str = 'auto'):
        self.backend_path = Path(backend_path)
        self.backend_type = backend_type
        self.backend = None

        # Detect backend type
        if backend_type == 'auto':
            if self.backend_path.suffix == '.vault':
                backend_type = 'hxv1'
            elif self.backend_path.is_dir():
                backend_type = 'npz'
            else:
                raise ValueError(f"Cannot auto-detect backend type for {backend_path}")

        # Initialize backend
        if backend_type == 'hxv1':
            self.backend = HXV1Backend(self.backend_path)
        elif backend_type == 'npz':
            self.backend = NPZBackend(self.backend_path)
        else:
            raise ValueError(f"Unknown backend type: {backend_type}")

        self.backend_type = backend_type

        # Virtual file names
        self.gguf_name = "model.gguf"
        self.gguf_path = f"/{self.gguf_name}"

        # Cache for reconstructed data
        self._cache = {}
        self._cache_max_bytes = 256 * 1024 * 1024  # 256MB
        self._cache_bytes = 0

        # Compute metadata
        now = time.time()
        self._st_mode_dir = stat.S_IFDIR | 0o555
        self._st_mode_file = stat.S_IFREG | 0o444
        self._mounted_at = now

        # Build metadata
        self.meta = {
            "version": "helix_cdc:v0.3.0-FROZEN-SUBSTRATE",
            "format": "HELIX_GGUF_FUSE.v1",
            "backend": backend_type,
            "path": str(self.backend_path.absolute()),
            "mounted_at": int(now),
        }

        if backend_type == 'hxv1':
            self.meta["vault_K"] = self.backend.K
            self.meta["vault_D"] = self.backend.D
        elif backend_type == 'npz':
            self.meta["num_layers"] = self.backend.num_layers
            if self.backend.manifest:
                self.meta["source_gguf"] = self.backend.manifest.get('source_gguf', 'unknown')
                self.meta["compression_ratio"] = self.backend.manifest.get('summary', {}).get('compression_ratio', 0)

        # Estimate total size for virtual GGUF
        self._estimate_gguf_size()

    def _estimate_gguf_size(self):
        """Estimate the virtual GGUF file size."""
        if self.backend_type == 'npz':
            manifest = self.backend.manifest
            if manifest and 'summary' in manifest:
                # Use original bytes from manifest
                self.total_size = manifest['summary'].get('original_bytes', 0)
            else:
                # Estimate from layer count
                self.total_size = self.backend.num_layers * 64 * 1024 * 1024  # Rough estimate
        elif self.backend_type == 'hxv1':
            # Estimate from vault dimensions
            self.total_size = self.backend.K * self.backend.D * 4  # float32
        else:
            self.total_size = 0

        self.meta["estimated_size"] = self.total_size

    def _split(self, path: str) -> List[str]:
        return [p for p in path.strip("/").split("/") if p]

    def _is_dir(self, parts: List[str]) -> bool:
        if len(parts) == 0:
            return True
        if parts[0] == "layers":
            return len(parts) == 1
        return False

    def _is_file(self, parts: List[str]) -> bool:
        if parts == ["meta.json"]:
            return True
        if parts == [self.gguf_name]:
            return True
        if len(parts) == 2 and parts[0] == "layers":
            return True
        return False

    # FUSE Operations
    def getattr(self, path, fh=None):
        parts = self._split(path)
        now = time.time()

        if path == self.gguf_path:
            return dict(
                st_mode=(stat.S_IFREG | 0o444),
                st_nlink=1,
                st_size=self.total_size,
                st_ctime=now, st_mtime=now, st_atime=now
            )

        if self._is_dir(parts):
            return dict(st_mode=self._st_mode_dir, st_nlink=2,
                       st_ctime=now, st_mtime=now, st_atime=now)

        if self._is_file(parts):
            if parts == ["meta.json"]:
                size = len(json.dumps(self.meta, indent=2).encode())
            elif parts[0] == "layers":
                # Layer file size
                try:
                    idx = int(parts[1].split(".")[0])
                    if self.backend_type == 'npz':
                        layer_info = self.backend.get_layer_info(idx)
                        size = layer_info.get('file_bytes', 4096)
                    else:
                        size = self.backend.layer_dim * 4
                except (ValueError, IndexError):
                    raise FuseOSError(errno.ENOENT)
            else:
                size = 0

            return dict(st_mode=self._st_mode_file, st_nlink=1, st_size=size,
                       st_ctime=now, st_mtime=now, st_atime=now)

        raise FuseOSError(errno.ENOENT)

    def readdir(self, path, fh):
        parts = self._split(path)

        if len(parts) == 0:
            entries = [".", "..", "meta.json", "layers", self.gguf_name]
            return entries

        if parts == ["layers"]:
            entries = [".", ".."]
            if self.backend_type == 'npz':
                for i in range(self.backend.num_layers):
                    entries.append(f"{i:04d}.npz")
            elif self.backend_type == 'hxv1':
                for i in range(self.backend.num_layers):
                    entries.append(f"{i:04d}.npy")
            return entries

        raise FuseOSError(errno.ENOENT)

    def open(self, path, flags):
        parts = self._split(path)
        if not self._is_file(parts) and path != self.gguf_path:
            raise FuseOSError(errno.ENOENT)
        if (flags & os.O_RDWR) or (flags & os.O_WRONLY):
            raise FuseOSError(errno.EROFS)
        return 0

    def read(self, path, size, offset, fh):
        parts = self._split(path)

        if parts == ["meta.json"]:
            data = json.dumps(self.meta, indent=2).encode()
            return data[offset:offset+size]

        if path == self.gguf_path:
            # Virtual GGUF read - this would need full GGUF reconstruction
            # For now, return zeros as placeholder
            # Real implementation would reconstruct from SVD components
            return b'\x00' * min(size, max(0, self.total_size - offset))

        if len(parts) == 2 and parts[0] == "layers":
            try:
                idx = int(parts[1].split(".")[0])

                # Check cache
                cache_key = f"layer_{idx}"
                if cache_key not in self._cache:
                    if self.backend_type == 'npz':
                        tensor = self.backend.reconstruct_layer(idx)
                    elif self.backend_type == 'hxv1':
                        tensor = self.backend.get_layer(idx)
                    else:
                        raise FuseOSError(errno.EIO)

                    self._cache[cache_key] = tensor.tobytes()

                data = self._cache[cache_key]
                return data[offset:offset+size]

            except (ValueError, IndexError):
                raise FuseOSError(errno.ENOENT)
            except Exception:
                raise FuseOSError(errno.EIO)

        raise FuseOSError(errno.ENOENT)

    # Read-only filesystem
    def write(self, *a, **k): raise FuseOSError(errno.EROFS)
    def mkdir(self, *a, **k): raise FuseOSError(errno.EROFS)
    def rmdir(self, *a, **k): raise FuseOSError(errno.EROFS)
    def unlink(self, *a, **k): raise FuseOSError(errno.EROFS)
    def rename(self, *a, **k): raise FuseOSError(errno.EROFS)
    def chmod(self, *a, **k): raise FuseOSError(errno.EROFS)
    def chown(self, *a, **k): raise FuseOSError(errno.EROFS)
    def truncate(self, *a, **k): raise FuseOSError(errno.EROFS)


def main():
    if not HAVE_FUSE:
        print("Error: fusepy not installed. Install with: pip install fusepy")
        sys.exit(1)

    ap = argparse.ArgumentParser(description="Helix GGUF FUSE mount")
    ap.add_argument("--hxv1", help="Path to HXV1 vault file")
    ap.add_argument("--npz", help="Path to NPZ-encoded directory")
    ap.add_argument("--mount", required=True, help="Mount point")
    ap.add_argument("--foreground", action="store_true", help="Run in foreground")
    args = ap.parse_args()

    # Determine backend
    if args.hxv1:
        backend_path = args.hxv1
        backend_type = 'hxv1'
    elif args.npz:
        backend_path = args.npz
        backend_type = 'npz'
    else:
        ap.error("Must specify --hxv1 or --npz")

    if not Path(backend_path).exists():
        print(f"Error: {backend_path} not found")
        sys.exit(1)

    if not Path(args.mount).exists():
        print(f"Creating mount point: {args.mount}")
        Path(args.mount).mkdir(parents=True)

    print(f"Mounting {backend_path} ({backend_type}) at {args.mount}")

    fuse = FUSE(
        HelixGGUFFuse(backend_path, backend_type),
        args.mount,
        nothreads=True,
        foreground=args.foreground,
        ro=True,
        allow_other=False,
        fsname=f"helix:{Path(backend_path).name}",
        uid=os.getuid(),
        gid=os.getgid()
    )


if __name__ == "__main__":
    main()
