#!/usr/bin/env python3
"""
FibPi Zlib Wrapper - Random-Access Compression Engine

Provides random-access capability to zlib/gzip compressed data by splitting
at Fibonacci positions and compressing each chunk independently.

Architecture:
    Original Data (N bytes)
        |
    Split at Fibonacci positions
        |
    [Chunk_0] [Chunk_1] [Chunk_2] ... [Chunk_N]
        |         |         |           |
    zlib.compress() for each chunk independently
        |
    [SHA_0][zlib_0][SHA_1][zlib_1]...[SHA_N][zlib_N]
        |
    FibPi header + embedded index

Key Feature: Each chunk is an independent zlib stream, enabling O(1) access
to any byte range (only decompress overlapping chunks).

Usage:
    from fibpi_zlib_wrapper import FibPiZlibWrapper

    wrapper = FibPiZlibWrapper()

    # Compress
    compressed = wrapper.compress_chunked(data)

    # Full decompress
    original = wrapper.decompress(compressed)

    # Random access - only decompresses needed chunks
    slice_data = wrapper.decompress_range(compressed, start=1_000_000, end=2_000_000)
"""

import hashlib
import struct
import zlib
import json
from pathlib import Path
from datetime import datetime
from dataclasses import dataclass
from typing import List, Tuple, Optional, BinaryIO, Dict


# Format magic bytes
MAGIC = b"FIBPI_ZLIB_v1\x00"  # 14 bytes, null-terminated for easy C interop
INDEX_MAGIC = b"FIBPI_ZIDX_v1\x00"  # 14 bytes for sidecar index

# Format version (bump when format changes incompatibly)
FORMAT_VERSION = 1

# Compression flags
COMPRESS_NONE = 0x00
COMPRESS_ZLIB = 0x01
COMPRESS_GZIP = 0x02  # RFC 1952 compatible
COMPRESS_BROTLI = 0x03
COMPRESS_LZ4 = 0x04
COMPRESS_ZSTD = 0x05
COMPRESS_LZMA = 0x06

# Struct formats and sizes (computed once, never hardcode again)
# ChunkIndex: chunk_id(I) + comp_offset(Q) + comp_len(I) + decomp_offset(Q) + decomp_len(I) + sha256(32s)
CHUNKINDEX_FMT = '>I Q I Q I 32s'
CHUNKINDEX_SIZE = struct.calcsize(CHUNKINDEX_FMT)  # 60 bytes

# Header format v1:
#   original_size(Q=8) + num_chunks(I=4) + index_offset(Q=8) + comp_flag(B=1) +
#   format_version(H=2) + entry_size(H=2) + reserved(3s=3) = 28 bytes
# This embeds version + entry_size so future readers can detect format changes
HEADER_FMT = '>Q I Q B H H 3s'
HEADER_SIZE = struct.calcsize(HEADER_FMT)  # 28 bytes


def generate_fibonacci(max_val: int = 100_000_000) -> List[int]:
    """Generate Fibonacci sequence up to max_val."""
    fib = [1, 1]
    while fib[-1] < max_val:
        fib.append(fib[-1] + fib[-2])
    return fib


def get_fibonacci_chunk_sizes(data_len: int, min_chunk: int = 4096) -> List[int]:
    """
    Get chunk sizes based on Fibonacci sequence.

    For small data, uses minimal chunks.
    For large data, scales Fibonacci sequence to create reasonable chunk sizes.

    Returns list of chunk sizes that sum to data_len.
    """
    if data_len <= min_chunk:
        return [data_len] if data_len > 0 else []

    fib = generate_fibonacci(data_len)

    # Find scaling factor to make chunks reasonable (16KB - 4MB range)
    # Aim for chunks that are neither too small (overhead) nor too large (latency)
    target_avg_chunk = min(4 * 1024 * 1024, max(16 * 1024, data_len // 100))

    # Use Fibonacci ratios for chunk boundaries
    cumulative = 0
    sizes = []
    fib_idx = 0

    while cumulative < data_len:
        # Scale Fibonacci number to reasonable chunk size
        fib_val = fib[min(fib_idx, len(fib) - 1)]
        scale = max(1, target_avg_chunk // max(1, fib_val))
        chunk_size = min(fib_val * scale, data_len - cumulative)
        chunk_size = max(min_chunk, chunk_size)  # Minimum chunk size

        if cumulative + chunk_size > data_len:
            chunk_size = data_len - cumulative

        if chunk_size > 0:
            sizes.append(chunk_size)
            cumulative += chunk_size

        fib_idx += 1

        # Safety: if we've gone through all Fibonacci numbers, use fixed size
        if fib_idx > len(fib):
            remaining = data_len - cumulative
            while remaining > 0:
                cs = min(target_avg_chunk, remaining)
                sizes.append(cs)
                remaining -= cs
            break

    return sizes


def sha256_digest(data: bytes) -> bytes:
    """Compute SHA256 hash as raw bytes (32 bytes)."""
    return hashlib.sha256(data).digest()


def sha256_hex(data: bytes) -> str:
    """Compute SHA256 hash as hex string (64 chars)."""
    return hashlib.sha256(data).hexdigest()


def unpack_header(header_bytes: bytes) -> Tuple[int, int, int, int, int, int]:
    """
    Unpack and validate a FibPi header.

    Returns:
        (original_size, num_chunks, index_offset, comp_flag, format_version, entry_size)

    Raises:
        ValueError: If format version or entry size is incompatible
    """
    if len(header_bytes) < HEADER_SIZE:
        raise ValueError(f"Header too short: {len(header_bytes)} < {HEADER_SIZE}")

    (original_size, num_chunks, index_offset, comp_flag,
     format_version, entry_size, _reserved) = struct.unpack(HEADER_FMT, header_bytes)

    # Validate format version
    if format_version > FORMAT_VERSION:
        raise ValueError(
            f"Unsupported format version: {format_version} (max supported: {FORMAT_VERSION})"
        )

    # Validate entry size (allows detection of corrupt files)
    if entry_size != CHUNKINDEX_SIZE and entry_size != 0:
        # entry_size=0 for backward compat with files written before versioning
        raise ValueError(
            f"Entry size mismatch: file has {entry_size}, expected {CHUNKINDEX_SIZE}"
        )

    return original_size, num_chunks, index_offset, comp_flag, format_version, entry_size


@dataclass
class ChunkIndex:
    """Index entry for a compressed chunk."""
    chunk_id: int
    compressed_offset: int      # Byte offset in compressed file
    compressed_len: int         # Length of compressed data
    decompressed_offset: int    # Byte offset in original data
    decompressed_len: int       # Length of original chunk
    sha256: bytes               # SHA256 of decompressed chunk (32 bytes)

    def to_bytes(self) -> bytes:
        """Serialize to binary format (CHUNKINDEX_SIZE bytes)."""
        return struct.pack(
            CHUNKINDEX_FMT,
            self.chunk_id,
            self.compressed_offset,
            self.compressed_len,
            self.decompressed_offset,
            self.decompressed_len,
            self.sha256
        )

    @classmethod
    def from_bytes(cls, data: bytes) -> 'ChunkIndex':
        """Deserialize from binary format."""
        chunk_id, comp_off, comp_len, decomp_off, decomp_len, sha = struct.unpack(
            CHUNKINDEX_FMT, data[:CHUNKINDEX_SIZE]
        )
        return cls(
            chunk_id=chunk_id,
            compressed_offset=comp_off,
            compressed_len=comp_len,
            decompressed_offset=decomp_off,
            decompressed_len=decomp_len,
            sha256=sha
        )


class FibPiZlibWrapper:
    """
    Random-access Gzip-compatible compression using Fibonacci-indexed chunks.

    Each chunk is compressed independently with zlib, enabling O(1) access
    to any byte range without decompressing the entire file.
    """

    def __init__(self, compression_level: int = 9):
        """
        Initialize the wrapper.

        Args:
            compression_level: zlib compression level (1-9, default 9 for max)
        """
        self.compression_level = compression_level
        self._stats = {
            'chunks_compressed': 0,
            'chunks_decompressed': 0,
            'bytes_in': 0,
            'bytes_out': 0
        }

    def compress_chunked(self, data: bytes, level: int = None) -> bytes:
        """
        Compress data with independent zlib streams per Fibonacci chunk.

        Args:
            data: Raw data to compress
            level: Compression level (1-9), defaults to instance level

        Returns:
            Compressed bytes with embedded index for random access
        """
        if level is None:
            level = self.compression_level

        data_len = len(data)
        chunk_sizes = get_fibonacci_chunk_sizes(data_len)

        # Build output
        output = bytearray(MAGIC)

        # Reserve space for header (will fill in at end)
        # Header: [original_size:8][num_chunks:4][index_offset:8][compression_flag:1][reserved:7]
        header_offset = len(output)
        output.extend(b'\x00' * HEADER_SIZE)  # 28-byte header

        # Compress each chunk and build index
        index_entries: List[ChunkIndex] = []
        decompressed_offset = 0

        for chunk_id, chunk_size in enumerate(chunk_sizes):
            chunk_data = data[decompressed_offset:decompressed_offset + chunk_size]

            # Compute SHA256 of original chunk
            chunk_sha = sha256_digest(chunk_data)

            # Compress with zlib
            compressed = zlib.compress(chunk_data, level)

            # Record compressed offset (relative to start of chunk data area)
            compressed_offset = len(output) - len(MAGIC) - HEADER_SIZE

            # Build index entry
            entry = ChunkIndex(
                chunk_id=chunk_id,
                compressed_offset=compressed_offset,
                compressed_len=len(compressed),
                decompressed_offset=decompressed_offset,
                decompressed_len=chunk_size,
                sha256=chunk_sha
            )
            index_entries.append(entry)

            # Write: [SHA256:32][comp_len:4][compressed_data]
            output.extend(chunk_sha)
            output.extend(struct.pack('>I', len(compressed)))
            output.extend(compressed)

            decompressed_offset += chunk_size
            self._stats['chunks_compressed'] += 1

        # Write index at end
        index_offset = len(output)
        for entry in index_entries:
            output.extend(entry.to_bytes())

        # Go back and fill in header (includes version + entry_size for future-proofing)
        header = struct.pack(
            HEADER_FMT,
            data_len,                    # original_size: 8 bytes
            len(index_entries),          # num_chunks: 4 bytes
            index_offset,                # index_offset: 8 bytes
            COMPRESS_ZLIB,               # compression_flag: 1 byte
            FORMAT_VERSION,              # format_version: 2 bytes
            CHUNKINDEX_SIZE,             # entry_size: 2 bytes (allows reader to verify)
            b'\x00' * 3                  # reserved: 3 bytes
        )
        output[header_offset:header_offset + HEADER_SIZE] = header

        self._stats['bytes_in'] += data_len
        self._stats['bytes_out'] += len(output)

        return bytes(output)

    def decompress(self, compressed: bytes, verify: bool = True) -> bytes:
        """
        Fully decompress data.

        Args:
            compressed: FibPi-compressed data
            verify: If True, verify SHA256 of each chunk

        Returns:
            Original data
        """
        # Parse header
        if not compressed.startswith(MAGIC):
            raise ValueError(f"Invalid magic bytes - expected {MAGIC!r}")

        pos = len(MAGIC)
        header_bytes = compressed[pos:pos + HEADER_SIZE]
        original_size, num_chunks, index_offset, comp_flag, fmt_ver, entry_sz = unpack_header(header_bytes)
        pos += HEADER_SIZE

        # Read index
        index_entries = self._read_index(compressed, index_offset, num_chunks)

        # Decompress all chunks
        output = bytearray()
        chunk_data_start = len(MAGIC) + HEADER_SIZE

        for entry in index_entries:
            chunk_offset = chunk_data_start + entry.compressed_offset

            # Read SHA256 and compressed length
            chunk_sha = compressed[chunk_offset:chunk_offset + 32]
            comp_len = struct.unpack('>I', compressed[chunk_offset + 32:chunk_offset + 36])[0]

            # Read and decompress
            comp_data = compressed[chunk_offset + 36:chunk_offset + 36 + comp_len]
            decompressed = zlib.decompress(comp_data)

            # Verify if requested
            if verify:
                computed_sha = sha256_digest(decompressed)
                if computed_sha != chunk_sha:
                    raise ValueError(
                        f"SHA256 mismatch for chunk {entry.chunk_id}: "
                        f"expected {chunk_sha.hex()}, got {computed_sha.hex()}"
                    )

            output.extend(decompressed)
            self._stats['chunks_decompressed'] += 1

        return bytes(output)

    def decompress_range(
        self,
        compressed: bytes,
        start: int,
        end: int,
        verify: bool = True
    ) -> bytes:
        """
        Decompress only a specific byte range from the original data.

        This is the key feature: O(1) access to any byte range by only
        decompressing the chunks that overlap [start, end).

        Args:
            compressed: FibPi-compressed data
            start: Start byte offset (inclusive)
            end: End byte offset (exclusive)
            verify: If True, verify SHA256 of decompressed chunks

        Returns:
            Bytes from [start, end) of the original data
        """
        # Parse header
        if not compressed.startswith(MAGIC):
            raise ValueError(f"Invalid magic bytes")

        pos = len(MAGIC)
        header_bytes = compressed[pos:pos + HEADER_SIZE]
        original_size, num_chunks, index_offset, comp_flag, fmt_ver, entry_sz = unpack_header(header_bytes)

        # Validate range
        if start < 0 or end > original_size or start >= end:
            raise ValueError(f"Invalid range [{start}, {end}) for data of size {original_size}")

        # Read index
        index_entries = self._read_index(compressed, index_offset, num_chunks)

        # Find overlapping chunks
        result = bytearray()
        chunk_data_start = len(MAGIC) + HEADER_SIZE

        for entry in index_entries:
            chunk_start = entry.decompressed_offset
            chunk_end = chunk_start + entry.decompressed_len

            # Check if this chunk overlaps our range
            if chunk_end <= start:
                continue  # Before our range
            if chunk_start >= end:
                break     # Past our range (index is sorted)

            # This chunk overlaps - decompress it
            chunk_offset = chunk_data_start + entry.compressed_offset
            chunk_sha = compressed[chunk_offset:chunk_offset + 32]
            comp_len = struct.unpack('>I', compressed[chunk_offset + 32:chunk_offset + 36])[0]

            comp_data = compressed[chunk_offset + 36:chunk_offset + 36 + comp_len]
            decompressed = zlib.decompress(comp_data)

            if verify:
                computed_sha = sha256_digest(decompressed)
                if computed_sha != chunk_sha:
                    raise ValueError(f"SHA256 mismatch for chunk {entry.chunk_id}")

            # Calculate slice within this chunk
            slice_start = max(0, start - chunk_start)
            slice_end = min(len(decompressed), end - chunk_start)

            result.extend(decompressed[slice_start:slice_end])
            self._stats['chunks_decompressed'] += 1

        return bytes(result)

    def _read_index(
        self,
        compressed: bytes,
        index_offset: int,
        num_chunks: int
    ) -> List[ChunkIndex]:
        """Read chunk index from compressed data."""
        entries = []
        pos = index_offset

        for _ in range(num_chunks):
            entry = ChunkIndex.from_bytes(compressed[pos:pos + CHUNKINDEX_SIZE])
            entries.append(entry)
            pos += CHUNKINDEX_SIZE

        return entries

    def get_stats(self) -> dict:
        """Get compression/decompression statistics."""
        stats = dict(self._stats)
        if stats['bytes_in'] > 0:
            stats['compression_ratio'] = stats['bytes_out'] / stats['bytes_in']
        return stats

    def compress_file(self, input_path: Path, output_path: Path) -> dict:
        """
        Compress a file to FibPi-zlib format.

        Args:
            input_path: Source file
            output_path: Destination file (.fibpi.gz)

        Returns:
            Stats dict with sizes and ratio
        """
        input_path = Path(input_path)
        output_path = Path(output_path)

        with open(input_path, 'rb') as f:
            data = f.read()

        compressed = self.compress_chunked(data)

        with open(output_path, 'wb') as f:
            f.write(compressed)

        original_sha = sha256_hex(data)

        return {
            'input': str(input_path),
            'output': str(output_path),
            'original_size': len(data),
            'compressed_size': len(compressed),
            'ratio': len(compressed) / len(data) if len(data) > 0 else 0,
            'original_sha256': original_sha,
            'num_chunks': self._stats['chunks_compressed']
        }

    def decompress_file(
        self,
        input_path: Path,
        output_path: Path,
        verify: bool = True
    ) -> dict:
        """
        Decompress a FibPi-zlib file.

        Args:
            input_path: Compressed file (.fibpi.gz)
            output_path: Destination file
            verify: Verify SHA256 checksums

        Returns:
            Stats dict
        """
        input_path = Path(input_path)
        output_path = Path(output_path)

        with open(input_path, 'rb') as f:
            compressed = f.read()

        data = self.decompress(compressed, verify=verify)

        with open(output_path, 'wb') as f:
            f.write(data)

        return {
            'input': str(input_path),
            'output': str(output_path),
            'compressed_size': len(compressed),
            'decompressed_size': len(data),
            'sha256': sha256_hex(data)
        }

    def extract_range_from_file(
        self,
        input_path: Path,
        start: int,
        end: int,
        verify: bool = True
    ) -> bytes:
        """
        Extract a byte range from a compressed file without full decompression.

        Args:
            input_path: Compressed file
            start: Start byte offset
            end: End byte offset
            verify: Verify SHA256

        Returns:
            Extracted bytes
        """
        with open(input_path, 'rb') as f:
            compressed = f.read()

        return self.decompress_range(compressed, start, end, verify=verify)


class FibPiZlibIndex:
    """
    In-memory index for fast random access to FibPi-compressed files.

    Usage:
        index = FibPiZlibIndex.from_file("large_file.fibpi.gz")
        chunk_data = index.get_chunk(5)
        byte_range = index.get_range(1_000_000, 2_000_000)
    """

    def __init__(
        self,
        filepath: Path,
        original_size: int,
        entries: List[ChunkIndex],
        chunk_data_offset: int
    ):
        self.filepath = Path(filepath)
        self.original_size = original_size
        self.entries = entries
        self.chunk_data_offset = chunk_data_offset
        self._file_handle: Optional[BinaryIO] = None

    @classmethod
    def from_file(cls, filepath: Path) -> 'FibPiZlibIndex':
        """Build index from a FibPi-compressed file."""
        filepath = Path(filepath)

        with open(filepath, 'rb') as f:
            magic = f.read(len(MAGIC))
            if magic != MAGIC:
                raise ValueError("Invalid FibPi file")

            header = f.read(HEADER_SIZE)
            original_size, num_chunks, index_offset, comp_flag, fmt_ver, entry_sz = unpack_header(header)

            # Read index
            f.seek(index_offset)
            entries = []
            for _ in range(num_chunks):
                entry = ChunkIndex.from_bytes(f.read(CHUNKINDEX_SIZE))
                entries.append(entry)

        return cls(
            filepath=filepath,
            original_size=original_size,
            entries=entries,
            chunk_data_offset=len(MAGIC) + HEADER_SIZE
        )

    def get_chunk(self, chunk_id: int, verify: bool = True) -> bytes:
        """Get a specific chunk by ID."""
        if chunk_id < 0 or chunk_id >= len(self.entries):
            raise IndexError(f"Chunk {chunk_id} out of range (0-{len(self.entries)-1})")

        entry = self.entries[chunk_id]

        with open(self.filepath, 'rb') as f:
            offset = self.chunk_data_offset + entry.compressed_offset
            f.seek(offset)

            sha = f.read(32)
            comp_len = struct.unpack('>I', f.read(4))[0]
            comp_data = f.read(comp_len)

        decompressed = zlib.decompress(comp_data)

        if verify:
            if sha256_digest(decompressed) != sha:
                raise ValueError(f"SHA256 mismatch for chunk {chunk_id}")

        return decompressed

    def get_range(self, start: int, end: int, verify: bool = True) -> bytes:
        """Get a byte range, only reading necessary chunks."""
        if start < 0 or end > self.original_size or start >= end:
            raise ValueError(f"Invalid range [{start}, {end})")

        result = bytearray()

        for entry in self.entries:
            chunk_start = entry.decompressed_offset
            chunk_end = chunk_start + entry.decompressed_len

            if chunk_end <= start:
                continue
            if chunk_start >= end:
                break

            chunk_data = self.get_chunk(entry.chunk_id, verify=verify)

            slice_start = max(0, start - chunk_start)
            slice_end = min(len(chunk_data), end - chunk_start)

            result.extend(chunk_data[slice_start:slice_end])

        return bytes(result)

    def get_metadata(self) -> dict:
        """Get file metadata."""
        return {
            'filepath': str(self.filepath),
            'original_size': self.original_size,
            'num_chunks': len(self.entries),
            'chunk_sizes': [e.decompressed_len for e in self.entries]
        }

    def list_chunks(self) -> List[dict]:
        """List all chunks with their metadata."""
        return [
            {
                'id': e.chunk_id,
                'compressed_offset': e.compressed_offset,
                'compressed_len': e.compressed_len,
                'decompressed_offset': e.decompressed_offset,
                'decompressed_len': e.decompressed_len,
                'sha256': e.sha256.hex()
            }
            for e in self.entries
        ]


def create_sidecar_index(compressed_path: Path, index_path: Path = None) -> Path:
    """
    Create a standalone sidecar index file (.fibpi.zidx) from a compressed file.

    This allows the index to be stored/transmitted separately from the data.

    Args:
        compressed_path: Path to FibPi-compressed file
        index_path: Optional output path (defaults to compressed_path + .zidx)

    Returns:
        Path to created index file
    """
    compressed_path = Path(compressed_path)
    if index_path is None:
        index_path = compressed_path.with_suffix(compressed_path.suffix + '.zidx')
    else:
        index_path = Path(index_path)

    # Read compressed file header and index
    with open(compressed_path, 'rb') as f:
        magic = f.read(len(MAGIC))
        if magic != MAGIC:
            raise ValueError("Invalid FibPi file")

        header = f.read(HEADER_SIZE)
        original_size, num_chunks, index_offset, comp_flag, fmt_ver, entry_sz = unpack_header(header)

        f.seek(index_offset)
        index_data = f.read(num_chunks * CHUNKINDEX_SIZE)

    # Write sidecar index
    with open(index_path, 'wb') as f:
        f.write(INDEX_MAGIC)
        f.write(struct.pack('>Q I', original_size, num_chunks))
        f.write(index_data)

    return index_path


if __name__ == '__main__':
    # Quick self-test
    import os

    print("FibPi Zlib Wrapper - Self Test")
    print("=" * 50)

    # Create test data
    test_data = os.urandom(1_000_000)  # 1MB random data
    print(f"Test data: {len(test_data):,} bytes")

    wrapper = FibPiZlibWrapper()

    # Compress
    compressed = wrapper.compress_chunked(test_data)
    ratio = len(compressed) / len(test_data)
    print(f"Compressed: {len(compressed):,} bytes ({ratio:.2%})")

    # Full decompress
    decompressed = wrapper.decompress(compressed)
    assert decompressed == test_data, "Full decompress mismatch!"
    print("Full decompress: VERIFIED")

    # Random access - get bytes 500000-600000
    slice_data = wrapper.decompress_range(compressed, 500_000, 600_000)
    assert slice_data == test_data[500_000:600_000], "Range decompress mismatch!"
    print(f"Random access [500000:600000]: VERIFIED ({len(slice_data):,} bytes)")

    # Test index
    print("\nIndex test:")
    stats = wrapper.get_stats()
    print(f"  Chunks compressed: {stats['chunks_compressed']}")
    print(f"  Chunks decompressed: {stats['chunks_decompressed']}")

    print("\nAll tests passed!")
