#!/usr/bin/env python3
"""
hxz_fuse.py - FUSE mount for .hxz archives with on-demand regrow
================================================================

Exposes .hxz (Helix compressed) archives as a virtual filesystem
with transparent on-demand regeneration of files.

Layout:
  /<mount>/
    <original file structure from manifest>

Files marked with mode="regrow" are regenerated on-demand using
the regrow() function registry.

Usage:
  python -m helix_cdc.fuse.hxz_fuse /path/to/archive.hxz /mnt/hxz
"""

import os
import sys
import struct
import json
import mmap
import tempfile
import hashlib
import logging
from errno import ENOENT, EROFS, EIO
from stat import S_IFDIR, S_IFREG
from pathlib import Path
from typing import Dict, Any, Optional, Set
import numpy as np

# Try to import dependencies
try:
    import zstandard as zstd
    HAVE_ZSTD = True
except ImportError:
    HAVE_ZSTD = False

try:
    from fuse import FUSE, FuseOSError, Operations, LoggingMixIn
    HAVE_FUSE = True
except ImportError:
    HAVE_FUSE = False
    class Operations:
        pass
    class LoggingMixIn:
        pass
    class FuseOSError(Exception):
        def __init__(self, errno):
            self.errno = errno

# Local regrow registry
REGROW_REGISTRY = {}


def register_regrow(name: str):
    """Decorator to register regrow functions."""
    def decorator(func):
        REGROW_REGISTRY[name] = func
        return func
    return decorator


def can_regrow(regrow_id: str) -> bool:
    """Check if a regrow function exists."""
    return regrow_id in REGROW_REGISTRY


def regrow(regrow_id: str, seed: Any, file_info: Dict[str, Any]) -> bytes:
    """Regrow a file using its registered function."""
    if regrow_id not in REGROW_REGISTRY:
        raise ValueError(f"Unknown regrow_id: {regrow_id}")
    return REGROW_REGISTRY[regrow_id](seed, file_info)


# Built-in regrow functions
@register_regrow('py_bytecode')
def grow_py_bytecode(seed: Any, file_info: Dict[str, Any]) -> bytes:
    """Regenerate Python bytecode from seed."""
    import hashlib
    h = hashlib.sha256(str(seed).encode()).digest()
    size = file_info.get("size", 512)
    # Deterministic pseudo-random bytecode using numpy RNG
    rng = np.random.default_rng(int.from_bytes(h[:8], 'little'))
    return rng.bytes(size)


@register_regrow('json_boot')
def grow_json_boot(seed: Any, file_info: Dict[str, Any]) -> bytes:
    """Regenerate JSON bootstrap from seed."""
    return json.dumps({
        "seed": str(seed),
        "size": file_info.get("size", 0),
        "regenerated": True
    }, indent=2).encode()


@register_regrow('tmpfs_volatile')
def grow_tmpfs_volatile(seed: Any, file_info: Dict[str, Any]) -> bytes:
    """Regenerate volatile tmpfs content."""
    import hashlib
    h = hashlib.sha256(str(seed).encode()).digest()
    size = file_info.get("size", 4096)
    # Deterministic pseudo-random content
    rng = np.random.default_rng(int.from_bytes(h[:8], 'little'))
    return rng.bytes(size)


@register_regrow('docs_md')
def grow_docs_md(seed: Any, file_info: Dict[str, Any]) -> bytes:
    """Regenerate markdown documentation."""
    return f"# {seed}\n\nRegenerated documentation file.\n".encode()


@register_regrow('bench_results')
def grow_bench_results(seed: Any, file_info: Dict[str, Any]) -> bytes:
    """Regenerate benchmark results."""
    return json.dumps({
        "benchmark": str(seed),
        "status": "regenerated",
        "timestamp": "on-demand"
    }, indent=2).encode()


class HelixHXZFuse(LoggingMixIn, Operations):
    """FUSE filesystem for .hxz archives with regrow support."""

    def __init__(self, hxz_path: str, cache_dir: Optional[str] = None):
        self.hxz_path = Path(hxz_path)
        self.cache_dir = Path(cache_dir) if cache_dir else Path(tempfile.mkdtemp(prefix="helix_cache_"))

        # Memory-map the .hxz file
        self.hxz_file = open(self.hxz_path, "rb")
        self.hxz_mmap = mmap.mmap(self.hxz_file.fileno(), 0, access=mmap.ACCESS_READ)

        # Parse header
        self._parse_header()

        # Load manifest
        self._load_manifest()

        # Track regrown files
        self.regrown: Dict[str, bytes] = {}

        # Statistics tracking
        self.stats = {
            "regrown_bytes": 0,
            "stored_bytes": 0,
            "regrow_calls": 0,
            "cache_hits": 0
        }

        # Track unique pages read
        self.read_pages: Dict[str, Set[int]] = {}

        print(f"Mounted {hxz_path}")
        print(f"  Files: {len(self.manifest.get('files', {}))}")
        print(f"  Cache: {self.cache_dir}")

    def _parse_header(self):
        """Parse .hxz header."""
        magic = self.hxz_mmap[0:4]
        if magic != b"HXZ\0":
            raise ValueError(f"Invalid magic: {magic}")

        version, = struct.unpack("<H", self.hxz_mmap[4:6])
        self.dict_sha256 = self.hxz_mmap[6:38]

        (self.manifest_ofs, self.manifest_len,
         self.anchors_ofs, self.anchors_len,
         self.payload_ofs) = struct.unpack("<QQQQQ", self.hxz_mmap[38:78])

        self.version = version

    def _load_manifest(self):
        """Load and decompress manifest."""
        if not HAVE_ZSTD:
            # Fallback: try to read as JSON directly
            try:
                manifest_data = self.hxz_mmap[self.manifest_ofs:self.manifest_ofs + self.manifest_len]
                self.manifest = json.loads(manifest_data)
                return
            except:
                self.manifest = {"files": {}}
                return

        manifest_compressed = self.hxz_mmap[self.manifest_ofs:self.manifest_ofs + self.manifest_len]

        dctx = zstd.ZstdDecompressor()
        manifest_json = dctx.decompress(manifest_compressed)

        self.manifest = json.loads(manifest_json)

        # Extract helix parameters
        helix_info = self.manifest.get("helix", {})
        self.helix_dim = helix_info.get("dim", int(os.environ.get("HELIX_D", 72)))

    def _get_file_info(self, path: str) -> Optional[Dict[str, Any]]:
        """Get file info from manifest."""
        if path.startswith("/"):
            path = path[1:]
        return self.manifest.get("files", {}).get(path)

    def _regrow_file(self, path: str, file_info: Dict[str, Any]) -> bytes:
        """Regrow a file on demand."""
        if path in self.regrown:
            self.stats["cache_hits"] += 1
            return self.regrown[path]

        regrow_id = file_info.get("regrow_id")
        seed = file_info.get("seed", path)

        if not regrow_id or not can_regrow(regrow_id):
            raise FuseOSError(ENOENT)

        # Regrow the file
        print(f"Regrowing {path} with {regrow_id}")
        data = regrow(regrow_id, seed, file_info)

        # Update statistics
        self.stats["regrow_calls"] += 1
        self.stats["regrown_bytes"] += len(data)

        # Cache in memory
        self.regrown[path] = data

        # Cache on disk
        cache_path = self.cache_dir / "regrown" / path
        cache_path.parent.mkdir(parents=True, exist_ok=True)
        cache_path.write_bytes(data)

        return data

    def getattr(self, path, fh=None):
        """Get file attributes."""
        if path == "/":
            return {
                "st_mode": S_IFDIR | 0o755,
                "st_nlink": 2,
                "st_size": 4096,
                "st_uid": os.getuid(),
                "st_gid": os.getgid(),
            }

        file_info = self._get_file_info(path)
        if file_info:
            mode = S_IFREG | 0o644
            size = file_info.get("size", 0)

            if file_info.get("mode") == "regrow":
                size = file_info.get("size", 4096)

            return {
                "st_mode": mode,
                "st_nlink": 1,
                "st_size": size,
                "st_uid": os.getuid(),
                "st_gid": os.getgid(),
                "st_mtime": file_info.get("mtime", 0),
            }

        # Check if directory
        dir_path = path[1:] if path.startswith("/") else path
        for file_path in self.manifest.get("files", {}):
            if file_path.startswith(dir_path + "/"):
                return {
                    "st_mode": S_IFDIR | 0o755,
                    "st_nlink": 2,
                    "st_size": 4096,
                    "st_uid": os.getuid(),
                    "st_gid": os.getgid(),
                }

        raise FuseOSError(ENOENT)

    def readdir(self, path, fh):
        """List directory contents."""
        entries = ['.', '..']

        if path == "/":
            prefix = ""
        else:
            prefix = path[1:] if path.startswith("/") else path
            if not prefix.endswith("/"):
                prefix += "/"

        seen = set()
        for file_path in self.manifest.get("files", {}):
            if prefix and not file_path.startswith(prefix):
                continue

            if prefix:
                rel_path = file_path[len(prefix):]
            else:
                rel_path = file_path

            parts = rel_path.split("/")
            if parts[0] and parts[0] not in seen:
                entries.append(parts[0])
                seen.add(parts[0])

        return entries

    def open(self, path, flags):
        """Open a file."""
        file_info = self._get_file_info(path)
        if not file_info:
            raise FuseOSError(ENOENT)

        if (flags & os.O_RDWR) or (flags & os.O_WRONLY):
            raise FuseOSError(EROFS)

        return hash(path) & 0x7FFFFFFF

    def read(self, path, size, offset, fh):
        """Read file data."""
        file_info = self._get_file_info(path)
        if not file_info:
            raise FuseOSError(ENOENT)

        if file_info.get("mode") == "regrow":
            # Regrow on demand
            clean_path = path[1:] if path.startswith("/") else path
            data = self._regrow_file(clean_path, file_info)
        else:
            # Read from stored data (payload section)
            file_offset = file_info.get("offset", 0)
            file_size = file_info.get("size", 0)

            start = self.payload_ofs + file_offset + offset
            end = min(start + size, self.payload_ofs + file_offset + file_size)

            data = self.hxz_mmap[start:end]
            self.stats["stored_bytes"] += len(data)
            return data

        return data[offset:offset + size]

    def release(self, path, fh):
        """Release file handle."""
        return 0

    def destroy(self, path):
        """Cleanup on unmount."""
        if hasattr(self, "hxz_mmap"):
            self.hxz_mmap.close()
        if hasattr(self, "hxz_file"):
            self.hxz_file.close()

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


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

    if len(sys.argv) < 3:
        print("Usage: python -m helix_cdc.fuse.hxz_fuse <archive.hxz> <mountpoint>")
        sys.exit(1)

    hxz_path = sys.argv[1]
    mount_point = sys.argv[2]

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

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

    fs = HelixHXZFuse(hxz_path)

    print(f"Mounting {hxz_path} at {mount_point}")
    print("Use Ctrl+C to unmount")

    try:
        FUSE(fs, mount_point, foreground=True, allow_other=False)
    except KeyboardInterrupt:
        print("\nUnmounting...")


if __name__ == "__main__":
    main()
