#!/usr/bin/env python3
"""
FibPi Gzip-Compatible Mode

Produces files that standard `gunzip` can read, while also creating a
sidecar index file that enables random-access decompression with Helix tools.

Two modes:
1. Standard gzip output + .fibpi.idx sidecar (gunzip compatible)
2. Concatenated gzip streams with embedded sync markers (self-contained)

Usage:
    # Create gzip-compatible file with sidecar
    from fibpi_gzip_compat import GzipCompatibleFibPi

    compat = GzipCompatibleFibPi()
    compat.write_gzip_with_index(data, Path("file.gz"))
    # Creates: file.gz (gunzip compatible) + file.gz.fibpi.idx (for random access)

    # Random access using sidecar
    slice_data = compat.read_range_with_index(Path("file.gz"), 1_000_000, 2_000_000)
"""

import gzip
import hashlib
import struct
import zlib
import io
import os
from pathlib import Path
from datetime import datetime
from dataclasses import dataclass, field
from typing import List, Tuple, Optional, BinaryIO


# Sidecar index magic
SIDECAR_MAGIC = b"FIBPI_IDX_v1\x00\x00"  # 14 bytes

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

# Sync marker for concatenated streams (16 bytes)
SYNC_MARKER = b"\xff\xfeFIBPI_SYNC\xfe\xff"

# Struct formats and sizes (computed once, never hardcode again)
# SidecarEntry: chunk_id(I) + gz_offset(Q) + gz_len(I) + orig_offset(Q) + orig_len(I) + sha256(32s) + crc32(I)
SIDECAR_ENTRY_FMT = '>I Q I Q I 32s I'
SIDECAR_ENTRY_SIZE = struct.calcsize(SIDECAR_ENTRY_FMT)  # 64 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."""
    if data_len <= min_chunk:
        return [data_len] if data_len > 0 else []

    fib = generate_fibonacci(data_len)
    target_avg = min(4 * 1024 * 1024, max(16 * 1024, data_len // 100))

    cumulative = 0
    sizes = []
    fib_idx = 0

    while cumulative < data_len:
        fib_val = fib[min(fib_idx, len(fib) - 1)]
        scale = max(1, target_avg // max(1, fib_val))
        chunk_size = min(fib_val * scale, data_len - cumulative)
        chunk_size = max(min_chunk, 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
        if fib_idx > len(fib):
            remaining = data_len - cumulative
            while remaining > 0:
                cs = min(target_avg, remaining)
                sizes.append(cs)
                remaining -= cs
            break

    return sizes


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


@dataclass
class SidecarEntry:
    """Index entry for sidecar index file."""
    chunk_id: int
    gzip_offset: int      # Byte offset in .gz file where this chunk's gzip stream starts
    gzip_len: int         # Length of the gzip stream for this chunk
    original_offset: int  # Byte offset in original uncompressed data
    original_len: int     # Length of original chunk
    sha256: bytes         # SHA256 of original chunk (32 bytes)
    crc32: int            # CRC32 for gzip validation

    def to_bytes(self) -> bytes:
        """Serialize to binary (SIDECAR_ENTRY_SIZE bytes)."""
        return struct.pack(
            SIDECAR_ENTRY_FMT,
            self.chunk_id,
            self.gzip_offset,
            self.gzip_len,
            self.original_offset,
            self.original_len,
            self.sha256,
            self.crc32
        )

    @classmethod
    def from_bytes(cls, data: bytes) -> 'SidecarEntry':
        """Deserialize from binary."""
        (chunk_id, gz_off, gz_len, orig_off, orig_len,
         sha, crc) = struct.unpack(SIDECAR_ENTRY_FMT, data[:SIDECAR_ENTRY_SIZE])
        return cls(
            chunk_id=chunk_id,
            gzip_offset=gz_off,
            gzip_len=gz_len,
            original_offset=orig_off,
            original_len=orig_len,
            sha256=sha,
            crc32=crc
        )


@dataclass
class SidecarIndex:
    """Complete sidecar index structure."""
    original_size: int
    original_sha256: bytes
    num_chunks: int
    entries: List[SidecarEntry] = field(default_factory=list)
    timestamp: str = ""
    source_filename: str = ""

    def to_bytes(self) -> bytes:
        """Serialize complete index to binary."""
        output = bytearray(SIDECAR_MAGIC)

        # Header
        ts_bytes = self.timestamp.encode()[:64].ljust(64, b'\x00')
        fn_bytes = self.source_filename.encode()[:256].ljust(256, b'\x00')

        output.extend(struct.pack(
            '>Q 32s I',
            self.original_size,
            self.original_sha256,
            self.num_chunks
        ))
        output.extend(ts_bytes)
        output.extend(fn_bytes)

        # Entries
        for entry in self.entries:
            output.extend(entry.to_bytes())

        return bytes(output)

    @classmethod
    def from_bytes(cls, data: bytes) -> 'SidecarIndex':
        """Deserialize from binary."""
        if not data.startswith(SIDECAR_MAGIC):
            raise ValueError("Invalid sidecar index magic")

        pos = len(SIDECAR_MAGIC)
        orig_size, orig_sha, num_chunks = struct.unpack(
            '>Q 32s I', data[pos:pos + 44]
        )
        pos += 44

        timestamp = data[pos:pos + 64].rstrip(b'\x00').decode()
        pos += 64
        filename = data[pos:pos + 256].rstrip(b'\x00').decode()
        pos += 256

        entries = []
        for _ in range(num_chunks):
            entry = SidecarEntry.from_bytes(data[pos:pos + SIDECAR_ENTRY_SIZE])
            entries.append(entry)
            pos += SIDECAR_ENTRY_SIZE

        return cls(
            original_size=orig_size,
            original_sha256=orig_sha,
            num_chunks=num_chunks,
            entries=entries,
            timestamp=timestamp,
            source_filename=filename
        )


class GzipCompatibleFibPi:
    """
    Creates gzip-compatible output with sidecar for random access.

    Standard gunzip can decompress the output normally.
    With the sidecar index, Helix tools can do random access.
    """

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

        Args:
            compression_level: gzip compression level (1-9)
        """
        self.compression_level = compression_level

    def compress_to_gzip(self, data: bytes, mtime: float = None) -> bytes:
        """
        Compress data to standard gzip format (single stream).

        This is standard gzip - no random access capability.
        """
        if mtime is None:
            mtime = datetime.now().timestamp()

        buf = io.BytesIO()
        with gzip.GzipFile(
            fileobj=buf,
            mode='wb',
            compresslevel=self.compression_level,
            mtime=mtime
        ) as gz:
            gz.write(data)

        return buf.getvalue()

    def write_gzip_with_index(
        self,
        data: bytes,
        output_path: Path,
        source_filename: str = None,
        mtime: float = None,
        timestamp: str = None
    ) -> Tuple[Path, Path]:
        """
        Write gzip file + sidecar index for random access.

        The .gz file is 100% standard gzip - gunzip will work.
        The .fibpi.idx sidecar enables random-access with Helix tools.

        NOTE: This uses concatenated gzip streams - each chunk is a
        complete gzip stream. Standard gunzip handles this correctly
        (it decompresses all streams in sequence).

        Args:
            data: Data to compress
            output_path: Output path (should end in .gz)
            source_filename: Original filename for metadata
            mtime: Modification time for gzip header (None = now, 0 = deterministic)
            timestamp: Sidecar timestamp string (None = now, provide fixed value for determinism)

        Returns:
            Tuple of (gzip_path, sidecar_path)
        """
        output_path = Path(output_path)
        sidecar_path = output_path.with_suffix(output_path.suffix + '.fibpi.idx')

        if source_filename is None:
            source_filename = output_path.stem

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

        # Build gzip with concatenated streams
        gzip_output = bytearray()
        entries: List[SidecarEntry] = []
        original_offset = 0
        # Use provided mtime, or current time for normal use
        # Pass mtime=0 for deterministic output (KAT, tests)
        if mtime is None:
            mtime = datetime.now().timestamp()

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

            # Record position before this gzip stream
            gzip_offset = len(gzip_output)

            # Create gzip stream for this chunk
            buf = io.BytesIO()
            with gzip.GzipFile(
                fileobj=buf,
                mode='wb',
                compresslevel=self.compression_level,
                mtime=mtime
            ) as gz:
                gz.write(chunk_data)

            gzip_stream = buf.getvalue()
            gzip_output.extend(gzip_stream)

            # Create sidecar entry
            entry = SidecarEntry(
                chunk_id=chunk_id,
                gzip_offset=gzip_offset,
                gzip_len=len(gzip_stream),
                original_offset=original_offset,
                original_len=chunk_size,
                sha256=sha256_digest(chunk_data),
                crc32=zlib.crc32(chunk_data) & 0xffffffff
            )
            entries.append(entry)

            original_offset += chunk_size

        # Write gzip file
        with open(output_path, 'wb') as f:
            f.write(gzip_output)

        # Write sidecar index
        # Use provided timestamp or current UTC time
        # For determinism (KAT), pass a fixed timestamp string
        idx_timestamp = timestamp if timestamp is not None else datetime.utcnow().isoformat() + "Z"
        index = SidecarIndex(
            original_size=data_len,
            original_sha256=sha256_digest(data),
            num_chunks=len(entries),
            entries=entries,
            timestamp=idx_timestamp,
            source_filename=source_filename
        )

        with open(sidecar_path, 'wb') as f:
            f.write(index.to_bytes())

        return output_path, sidecar_path

    def read_range_with_index(
        self,
        gzip_path: Path,
        start: int,
        end: int,
        sidecar_path: Path = None,
        verify: bool = True
    ) -> bytes:
        """
        Read a byte range using the sidecar index for random access.

        Only decompresses the chunks that overlap [start, end).

        Args:
            gzip_path: Path to .gz file
            start: Start byte offset
            end: End byte offset
            sidecar_path: Path to sidecar (default: gzip_path + .fibpi.idx)
            verify: Verify SHA256 of decompressed chunks

        Returns:
            Bytes from [start, end) of original data
        """
        gzip_path = Path(gzip_path)
        if sidecar_path is None:
            sidecar_path = gzip_path.with_suffix(gzip_path.suffix + '.fibpi.idx')
        else:
            sidecar_path = Path(sidecar_path)

        # Load sidecar index
        with open(sidecar_path, 'rb') as f:
            index = SidecarIndex.from_bytes(f.read())

        if start < 0 or end > index.original_size or start >= end:
            raise ValueError(f"Invalid range [{start}, {end})")

        result = bytearray()

        with open(gzip_path, 'rb') as f:
            for entry in index.entries:
                chunk_start = entry.original_offset
                chunk_end = chunk_start + entry.original_len

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

                # Seek to this chunk's gzip stream
                f.seek(entry.gzip_offset)
                gzip_stream = f.read(entry.gzip_len)

                # Decompress
                decompressed = gzip.decompress(gzip_stream)

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

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

                result.extend(decompressed[slice_start:slice_end])

        return bytes(result)

    def full_decompress_with_index(
        self,
        gzip_path: Path,
        sidecar_path: Path = None,
        verify: bool = True
    ) -> bytes:
        """
        Fully decompress using sidecar index (verifies each chunk).

        Note: Standard `gunzip` also works, but won't verify chunks.
        """
        gzip_path = Path(gzip_path)
        if sidecar_path is None:
            sidecar_path = gzip_path.with_suffix(gzip_path.suffix + '.fibpi.idx')

        with open(sidecar_path, 'rb') as f:
            index = SidecarIndex.from_bytes(f.read())

        return self.read_range_with_index(
            gzip_path, 0, index.original_size, sidecar_path, verify
        )

    def verify_file(
        self,
        gzip_path: Path,
        sidecar_path: Path = None
    ) -> dict:
        """
        Verify a gzip file using its sidecar index.

        Returns verification status for each chunk.
        """
        gzip_path = Path(gzip_path)
        if sidecar_path is None:
            sidecar_path = gzip_path.with_suffix(gzip_path.suffix + '.fibpi.idx')

        with open(sidecar_path, 'rb') as f:
            index = SidecarIndex.from_bytes(f.read())

        results = []
        all_ok = True

        with open(gzip_path, 'rb') as f:
            for entry in index.entries:
                f.seek(entry.gzip_offset)
                gzip_stream = f.read(entry.gzip_len)

                try:
                    decompressed = gzip.decompress(gzip_stream)
                    sha_ok = sha256_digest(decompressed) == entry.sha256
                    crc_ok = (zlib.crc32(decompressed) & 0xffffffff) == entry.crc32
                    len_ok = len(decompressed) == entry.original_len

                    status = "OK" if (sha_ok and crc_ok and len_ok) else "MISMATCH"
                    if status != "OK":
                        all_ok = False

                    results.append({
                        'chunk_id': entry.chunk_id,
                        'status': status,
                        'sha256_ok': sha_ok,
                        'crc32_ok': crc_ok,
                        'length_ok': len_ok
                    })
                except Exception as e:
                    all_ok = False
                    results.append({
                        'chunk_id': entry.chunk_id,
                        'status': 'ERROR',
                        'error': str(e)
                    })

        return {
            'file': str(gzip_path),
            'sidecar': str(sidecar_path),
            'overall_status': 'OK' if all_ok else 'FAILED',
            'original_size': index.original_size,
            'num_chunks': index.num_chunks,
            'chunks': results
        }


class ConcatenatedGzipReader:
    """
    Reader for concatenated gzip streams without sidecar.

    Standard gzip files can contain multiple concatenated streams.
    This reader can iterate through them or access by index.

    Note: Without a sidecar, accessing stream N requires reading
    through streams 0..N-1 first (no true random access).
    """

    def __init__(self, filepath: Path):
        self.filepath = Path(filepath)
        self._stream_offsets: Optional[List[Tuple[int, int]]] = None

    def _scan_streams(self) -> List[Tuple[int, int]]:
        """
        Scan file to find all gzip stream boundaries.

        Returns list of (offset, length) for each stream.
        """
        if self._stream_offsets is not None:
            return self._stream_offsets

        offsets = []

        with open(self.filepath, 'rb') as f:
            while True:
                pos = f.tell()
                header = f.read(10)

                if len(header) < 10:
                    break

                # Check for gzip magic
                if header[:2] != b'\x1f\x8b':
                    break

                # Skip to end of this stream by decompressing
                f.seek(pos)
                try:
                    with gzip.GzipFile(fileobj=f) as gz:
                        while gz.read(1024 * 1024):
                            pass
                except:
                    break

                end_pos = f.tell()
                offsets.append((pos, end_pos - pos))

        self._stream_offsets = offsets
        return offsets

    def get_stream_count(self) -> int:
        """Get number of gzip streams in file."""
        return len(self._scan_streams())

    def get_stream(self, index: int) -> bytes:
        """Get decompressed data from a specific stream."""
        offsets = self._scan_streams()

        if index < 0 or index >= len(offsets):
            raise IndexError(f"Stream {index} out of range")

        offset, length = offsets[index]

        with open(self.filepath, 'rb') as f:
            f.seek(offset)
            with gzip.GzipFile(fileobj=f) as gz:
                return gz.read()

    def iter_streams(self):
        """Iterate through all decompressed streams."""
        for i in range(self.get_stream_count()):
            yield self.get_stream(i)


def create_sidecar_from_concatenated_gzip(
    gzip_path: Path,
    sidecar_path: Path = None
) -> Path:
    """
    Create a sidecar index from an existing concatenated gzip file.

    This allows adding random-access capability to files created
    without Helix tools (e.g., by concatenating gzip streams).

    Args:
        gzip_path: Path to concatenated gzip file
        sidecar_path: Output path (default: gzip_path + .fibpi.idx)

    Returns:
        Path to created sidecar
    """
    gzip_path = Path(gzip_path)
    if sidecar_path is None:
        sidecar_path = gzip_path.with_suffix(gzip_path.suffix + '.fibpi.idx')

    reader = ConcatenatedGzipReader(gzip_path)
    offsets = reader._scan_streams()

    entries = []
    original_offset = 0
    all_data = bytearray()

    with open(gzip_path, 'rb') as f:
        for chunk_id, (gz_offset, gz_len) in enumerate(offsets):
            f.seek(gz_offset)
            with gzip.GzipFile(fileobj=f) as gz:
                chunk_data = gz.read()

            all_data.extend(chunk_data)

            entry = SidecarEntry(
                chunk_id=chunk_id,
                gzip_offset=gz_offset,
                gzip_len=gz_len,
                original_offset=original_offset,
                original_len=len(chunk_data),
                sha256=sha256_digest(chunk_data),
                crc32=zlib.crc32(chunk_data) & 0xffffffff
            )
            entries.append(entry)

            original_offset += len(chunk_data)

    index = SidecarIndex(
        original_size=len(all_data),
        original_sha256=sha256_digest(bytes(all_data)),
        num_chunks=len(entries),
        entries=entries,
        timestamp=datetime.utcnow().isoformat() + "Z",
        source_filename=gzip_path.stem
    )

    with open(sidecar_path, 'wb') as f:
        f.write(index.to_bytes())

    return sidecar_path


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

    print("FibPi Gzip-Compatible Mode - Self Test")
    print("=" * 50)

    # Create test data
    test_data = os.urandom(500_000)  # 500KB
    print(f"Test data: {len(test_data):,} bytes")

    compat = GzipCompatibleFibPi()

    # Write with sidecar
    gz_path, idx_path = compat.write_gzip_with_index(
        test_data,
        Path("/tmp/test_fibpi_compat.gz")
    )
    print(f"Created: {gz_path}")
    print(f"Sidecar: {idx_path}")

    gz_size = gz_path.stat().st_size
    idx_size = idx_path.stat().st_size
    print(f"  .gz size: {gz_size:,} bytes ({gz_size/len(test_data):.2%})")
    print(f"  .idx size: {idx_size:,} bytes")

    # Test standard gunzip compatibility
    print("\nTesting standard gunzip compatibility...")
    import subprocess
    result = subprocess.run(
        ['gunzip', '-c', str(gz_path)],
        capture_output=True
    )
    gunzip_data = result.stdout
    assert gunzip_data == test_data, "gunzip output mismatch!"
    print("  Standard gunzip: COMPATIBLE")

    # Test random access
    print("\nTesting random access...")
    slice_data = compat.read_range_with_index(gz_path, 100_000, 200_000)
    assert slice_data == test_data[100_000:200_000], "Random access mismatch!"
    print(f"  Range [100000:200000]: VERIFIED ({len(slice_data):,} bytes)")

    # Test verification
    print("\nTesting verification...")
    verify_result = compat.verify_file(gz_path)
    print(f"  Overall: {verify_result['overall_status']}")
    print(f"  Chunks: {verify_result['num_chunks']}")

    # Cleanup
    gz_path.unlink()
    idx_path.unlink()

    print("\nAll tests passed!")
