#!/usr/bin/env python3
"""
FibPi Index Manager - Persistent Index/Manifest for FibPi SHA Anchor Files

Eliminates O(n) file scan on open by persisting chunk offsets to a sidecar .fibpi.idx file.

Index File Format:
    [FIBPI_IDX_v1]     12 bytes - magic
    [mtime]             8 bytes - uint64 microseconds
    [n_chunks]          4 bytes - uint32
    [chunk_0_offset]    8 bytes - uint64 (file offset to chunk start)
    [chunk_0_len]       4 bytes - uint32 (encoded chunk length)
    [chunk_0_anchor]   32 bytes - raw SHA256 bytes
    ... repeat for each chunk

Usage:
    from fibpi_index import FibPiIndexManager, FibPiRandomAccessIndexed

    # Create index for a FibPi file
    mgr = FibPiIndexManager()
    mgr.create_index(Path("data.fibpi"))

    # Fast random access using index
    ra = FibPiRandomAccessIndexed(Path("data.fibpi"))
    chunk = ra.get_chunk_by_index(5)
"""

import os
import json
import struct
import hashlib
from pathlib import Path
from typing import List, Optional, Tuple, NamedTuple
from dataclasses import dataclass

# Import base FibPi components - try relative, then absolute, then fallback
try:
    # Try relative import (package context)
    from .fibpi_sha_anchor import (
        HEADER_MAGIC, ANCHOR_PREFIX, ANCHOR_LEN, CHUNK_PREFIX, CHUNK_SUFFIX,
        FibPiRandomAccess, sha256_hex
    )
except ImportError:
    try:
        # Try absolute import (direct execution)
        from fibpi_sha_anchor import (
            HEADER_MAGIC, ANCHOR_PREFIX, ANCHOR_LEN, CHUNK_PREFIX, CHUNK_SUFFIX,
            FibPiRandomAccess, sha256_hex
        )
    except ImportError:
        # Fallback for standalone testing
        HEADER_MAGIC = b"FIBPI3D_SHA_v1"
        ANCHOR_PREFIX = b"::SHA::"
        ANCHOR_LEN = 64
        CHUNK_PREFIX = b"::B64::"
        CHUNK_SUFFIX = b"::END::"

        def sha256_hex(data: bytes) -> str:
            return hashlib.sha256(data).hexdigest()


# Index format constants
INDEX_MAGIC = b"FIBPI_IDX_v1"  # 12 bytes
INDEX_VERSION = 1
CHUNK_ENTRY_SIZE = 8 + 4 + 32  # offset(8) + length(4) + anchor(32) = 44 bytes


@dataclass
class ChunkInfo:
    """Information about a single chunk in the index."""
    offset: int       # File offset to start of chunk entry (at ANCHOR_PREFIX)
    length: int       # Length of entire chunk entry in file
    anchor: bytes     # Raw SHA256 bytes (32 bytes)

    @property
    def anchor_hex(self) -> str:
        """Get anchor as hex string."""
        return self.anchor.hex()


class FibPiIndexManager:
    """Manages persistent .fibpi.idx sidecar files."""

    def __init__(self):
        self.stats = {
            'indexes_created': 0,
            'indexes_loaded': 0,
            'stale_invalidated': 0
        }

    def get_index_path(self, fibpi_path: Path) -> Path:
        """Get the sidecar index path for a FibPi file."""
        return fibpi_path.parent / (fibpi_path.name + ".idx")

    def create_index(self, fibpi_path: Path) -> Path:
        """
        Create sidecar index from FibPi file.

        Structure: [MAGIC][mtime][n_chunks][(offset, len, anchor) * n_chunks]
        """
        fibpi_path = Path(fibpi_path)
        idx_path = self.get_index_path(fibpi_path)

        # Get file mtime
        stat = fibpi_path.stat()
        mtime_us = int(stat.st_mtime * 1_000_000)  # microseconds

        # Read and index the FibPi file
        with open(fibpi_path, 'rb') as f:
            data = f.read()

        if not data.startswith(HEADER_MAGIC):
            raise ValueError(f"Invalid FibPi file: {fibpi_path}")

        # Parse to find chunk positions
        chunk_infos = self._parse_chunks(data)

        # Write index file
        with open(idx_path, 'wb') as f:
            # Magic (12 bytes)
            f.write(INDEX_MAGIC)

            # mtime (8 bytes, uint64)
            f.write(struct.pack('>Q', mtime_us))

            # n_chunks (4 bytes, uint32)
            f.write(struct.pack('>I', len(chunk_infos)))

            # Chunk entries
            for chunk in chunk_infos:
                f.write(struct.pack('>Q', chunk.offset))   # offset (8 bytes)
                f.write(struct.pack('>I', chunk.length))   # length (4 bytes)
                f.write(chunk.anchor)                       # anchor (32 bytes raw)

        self.stats['indexes_created'] += 1
        return idx_path

    def _parse_chunks(self, data: bytes) -> List[ChunkInfo]:
        """Parse FibPi file to extract chunk positions and anchors."""
        chunks = []

        # Skip header and metadata
        pos = len(HEADER_MAGIC)

        # Read metadata length
        meta_len = int.from_bytes(data[pos:pos+4], 'big')
        pos += 4 + meta_len

        # Parse chunks
        while pos < len(data):
            if data[pos:pos+len(ANCHOR_PREFIX)] != ANCHOR_PREFIX:
                break

            chunk_start = pos
            pos += len(ANCHOR_PREFIX)

            # Read anchor (64 hex chars -> 32 bytes raw)
            anchor_hex = data[pos:pos+ANCHOR_LEN].decode()
            anchor_raw = bytes.fromhex(anchor_hex)
            pos += ANCHOR_LEN

            # Skip to end of chunk
            suffix_pos = data.find(CHUNK_SUFFIX, pos)
            if suffix_pos == -1:
                break
            pos = suffix_pos + len(CHUNK_SUFFIX)

            chunk_len = pos - chunk_start

            chunks.append(ChunkInfo(
                offset=chunk_start,
                length=chunk_len,
                anchor=anchor_raw
            ))

        return chunks

    def load_index(self, fibpi_path: Path) -> Optional[List[ChunkInfo]]:
        """
        Load index if valid (mtime matches).

        Returns None if index doesn't exist or is stale.
        """
        fibpi_path = Path(fibpi_path)
        idx_path = self.get_index_path(fibpi_path)

        if not idx_path.exists():
            return None

        # Check mtime match
        try:
            stat = fibpi_path.stat()
            current_mtime_us = int(stat.st_mtime * 1_000_000)

            with open(idx_path, 'rb') as f:
                # Verify magic
                magic = f.read(12)
                if magic != INDEX_MAGIC:
                    return None

                # Read stored mtime
                stored_mtime_us = struct.unpack('>Q', f.read(8))[0]

                # Check if stale
                if stored_mtime_us != current_mtime_us:
                    self.stats['stale_invalidated'] += 1
                    return None

                # Read chunk count
                n_chunks = struct.unpack('>I', f.read(4))[0]

                # Read chunk entries
                chunks = []
                for _ in range(n_chunks):
                    offset = struct.unpack('>Q', f.read(8))[0]
                    length = struct.unpack('>I', f.read(4))[0]
                    anchor = f.read(32)

                    chunks.append(ChunkInfo(
                        offset=offset,
                        length=length,
                        anchor=anchor
                    ))

            self.stats['indexes_loaded'] += 1
            return chunks

        except Exception as e:
            # Index corrupted or unreadable
            return None

    def invalidate_if_stale(self, fibpi_path: Path) -> bool:
        """
        Check mtime and delete stale index.

        Returns True if index was deleted (stale), False if valid or doesn't exist.
        """
        fibpi_path = Path(fibpi_path)
        idx_path = self.get_index_path(fibpi_path)

        if not idx_path.exists():
            return False

        try:
            stat = fibpi_path.stat()
            current_mtime_us = int(stat.st_mtime * 1_000_000)

            with open(idx_path, 'rb') as f:
                f.read(12)  # Skip magic
                stored_mtime_us = struct.unpack('>Q', f.read(8))[0]

            if stored_mtime_us != current_mtime_us:
                idx_path.unlink()
                self.stats['stale_invalidated'] += 1
                return True

            return False

        except Exception:
            # If we can't check, delete to be safe
            try:
                idx_path.unlink()
            except Exception:
                pass
            return True

    def get_index_stats(self, fibpi_path: Path) -> Optional[dict]:
        """Get stats about an index file."""
        idx_path = self.get_index_path(fibpi_path)

        if not idx_path.exists():
            return None

        try:
            stat = idx_path.stat()
            with open(idx_path, 'rb') as f:
                magic = f.read(12)
                if magic != INDEX_MAGIC:
                    return None

                mtime_us = struct.unpack('>Q', f.read(8))[0]
                n_chunks = struct.unpack('>I', f.read(4))[0]

            return {
                'index_path': str(idx_path),
                'index_size': stat.st_size,
                'stored_mtime': mtime_us / 1_000_000,
                'n_chunks': n_chunks,
                'entry_size': CHUNK_ENTRY_SIZE,
                'header_size': 12 + 8 + 4
            }
        except Exception:
            return None


class FibPiRandomAccessIndexed:
    """
    Random access to FibPi SHA anchor encoded files using persistent index.

    Optimized version of FibPiRandomAccess that uses sidecar .idx files
    for O(1) open time instead of O(n) file scanning.
    """

    def __init__(self, filepath: Path, use_index: bool = True):
        """
        Initialize random access with optional index acceleration.

        Args:
            filepath: Path to .fibpi file
            use_index: Whether to use/create persistent index
        """
        self.filepath = Path(filepath)
        self.meta = None
        self.chunk_offsets: List[ChunkInfo] = []
        self._index_manager = FibPiIndexManager()

        if use_index:
            # Try to load from index
            cached = self._index_manager.load_index(filepath)
            if cached:
                self.chunk_offsets = cached
                self._load_metadata_only()
                return

        # Fallback to O(n) scan
        self._index_file()

        # Cache for next time
        if use_index:
            try:
                self._index_manager.create_index(filepath)
            except Exception:
                pass  # Non-fatal if index creation fails

    def _load_metadata_only(self):
        """Load just the metadata without scanning chunks."""
        with open(self.filepath, 'rb') as f:
            # Verify header
            magic = f.read(len(HEADER_MAGIC))
            if magic != HEADER_MAGIC:
                raise ValueError("Invalid FibPi file")

            # Read metadata
            meta_len = int.from_bytes(f.read(4), 'big')
            meta_json = f.read(meta_len)
            self.meta = json.loads(meta_json.decode())

    def _index_file(self):
        """Build index of chunk positions for random access (fallback)."""
        with open(self.filepath, 'rb') as f:
            data = f.read()

        if not data.startswith(HEADER_MAGIC):
            raise ValueError("Invalid FibPi file")

        pos = len(HEADER_MAGIC)

        # Read metadata
        meta_len = int.from_bytes(data[pos:pos+4], 'big')
        pos += 4
        self.meta = json.loads(data[pos:pos+meta_len].decode())
        pos += meta_len

        # Index all chunks
        while pos < len(data):
            if data[pos:pos+len(ANCHOR_PREFIX)] != ANCHOR_PREFIX:
                break

            chunk_start = pos
            pos += len(ANCHOR_PREFIX)

            # Read anchor
            anchor_hex = data[pos:pos+ANCHOR_LEN].decode()
            anchor_raw = bytes.fromhex(anchor_hex)
            pos += ANCHOR_LEN

            # Skip to end of chunk
            suffix_pos = data.find(CHUNK_SUFFIX, pos)
            if suffix_pos == -1:
                break
            pos = suffix_pos + len(CHUNK_SUFFIX)

            self.chunk_offsets.append(ChunkInfo(
                offset=chunk_start,
                length=pos - chunk_start,
                anchor=anchor_raw
            ))

    def get_chunk_count(self) -> int:
        """Get total number of chunks."""
        return len(self.chunk_offsets)

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

        info = self.chunk_offsets[idx]

        with open(self.filepath, 'rb') as f:
            f.seek(info.offset + len(ANCHOR_PREFIX) + ANCHOR_LEN + len(CHUNK_PREFIX))

            # Read until CHUNK_SUFFIX
            # We know the exact length from the index
            remaining = info.length - len(ANCHOR_PREFIX) - ANCHOR_LEN - len(CHUNK_PREFIX) - len(CHUNK_SUFFIX)
            b64_chunk = f.read(remaining)

        # Decode base64
        import base64
        decoded = base64.b64decode(b64_chunk)

        # Verify if requested
        if verify:
            computed = sha256_hex(decoded)
            if computed != info.anchor_hex:
                raise ValueError(f"Anchor mismatch for chunk {idx}")

        return decoded

    def get_chunk_by_anchor(self, anchor: str) -> Optional[bytes]:
        """Get a chunk by its SHA256 anchor (hex string, full or prefix)."""
        anchor_lower = anchor.lower()
        for idx, info in enumerate(self.chunk_offsets):
            if info.anchor_hex == anchor_lower or info.anchor_hex.startswith(anchor_lower):
                return self.get_chunk_by_index(idx)
        return None

    def search_pattern(self, pattern: bytes) -> List[Tuple[int, int]]:
        """Search for a pattern across all chunks."""
        matches = []
        for idx in range(len(self.chunk_offsets)):
            chunk = self.get_chunk_by_index(idx, verify=False)
            offset = 0
            while True:
                pos = chunk.find(pattern, offset)
                if pos == -1:
                    break
                matches.append((idx, pos))
                offset = pos + 1
        return matches

    def get_byte_range(self, start: int, end: int) -> bytes:
        """Get a byte range from the original data."""
        positions = self.meta.get('fibonacci_positions', [])
        result = bytearray()

        prev_pos = 0
        for idx, pos in enumerate(positions):
            if idx >= len(self.chunk_offsets):
                break

            chunk_start = prev_pos
            chunk_end = pos

            if chunk_end > start and chunk_start < end:
                chunk = self.get_chunk_by_index(idx, verify=False)
                slice_start = max(0, start - chunk_start)
                slice_end = min(len(chunk), end - chunk_start)
                result.extend(chunk[slice_start:slice_end])

            prev_pos = pos
            if prev_pos >= end:
                break

        return bytes(result)

    def list_anchors(self) -> List[str]:
        """List all SHA256 anchors in order."""
        return [info.anchor_hex for info in self.chunk_offsets]

    def get_metadata(self) -> dict:
        """Get file metadata."""
        return self.meta

    def get_index_info(self) -> dict:
        """Get information about the index being used."""
        return {
            'chunk_count': len(self.chunk_offsets),
            'indexed': bool(self._index_manager.load_index(self.filepath)),
            'index_path': str(self._index_manager.get_index_path(self.filepath))
        }


def benchmark_indexed_vs_unindexed(fibpi_path: Path) -> dict:
    """Benchmark indexed vs non-indexed random access open time."""
    import time

    fibpi_path = Path(fibpi_path)

    # Ensure index exists
    mgr = FibPiIndexManager()
    idx_path = mgr.get_index_path(fibpi_path)
    if not idx_path.exists():
        mgr.create_index(fibpi_path)

    # Benchmark indexed
    t0 = time.perf_counter()
    ra_indexed = FibPiRandomAccessIndexed(fibpi_path, use_index=True)
    indexed_time = time.perf_counter() - t0

    # Benchmark non-indexed
    t0 = time.perf_counter()
    ra_scan = FibPiRandomAccessIndexed(fibpi_path, use_index=False)
    scan_time = time.perf_counter() - t0

    return {
        'indexed_open_ms': indexed_time * 1000,
        'scan_open_ms': scan_time * 1000,
        'speedup': scan_time / indexed_time if indexed_time > 0 else float('inf'),
        'chunk_count': ra_indexed.get_chunk_count(),
        'index_size': idx_path.stat().st_size if idx_path.exists() else 0
    }


def main():
    import argparse

    parser = argparse.ArgumentParser(description='FibPi Index Manager')
    parser.add_argument('command', choices=['create', 'check', 'benchmark', 'stats'],
                        help='Command to run')
    parser.add_argument('file', help='FibPi file path')

    args = parser.parse_args()
    fibpi_path = Path(args.file)

    if not fibpi_path.exists():
        print(f"Error: File not found: {fibpi_path}")
        return 1

    mgr = FibPiIndexManager()

    if args.command == 'create':
        idx_path = mgr.create_index(fibpi_path)
        stats = mgr.get_index_stats(fibpi_path)
        print(f"Index created: {idx_path}")
        print(f"  Chunks indexed: {stats['n_chunks']}")
        print(f"  Index size: {stats['index_size']:,} bytes")

    elif args.command == 'check':
        cached = mgr.load_index(fibpi_path)
        if cached:
            print(f"Index VALID for {fibpi_path}")
            print(f"  Chunks: {len(cached)}")
        else:
            print(f"Index STALE or MISSING for {fibpi_path}")
            print("  Run 'create' to rebuild")

    elif args.command == 'benchmark':
        results = benchmark_indexed_vs_unindexed(fibpi_path)
        print(f"Benchmark results for {fibpi_path}:")
        print(f"  Indexed open:  {results['indexed_open_ms']:.3f} ms")
        print(f"  Scan open:     {results['scan_open_ms']:.3f} ms")
        print(f"  Speedup:       {results['speedup']:.1f}x")
        print(f"  Chunks:        {results['chunk_count']}")
        print(f"  Index size:    {results['index_size']:,} bytes")

    elif args.command == 'stats':
        stats = mgr.get_index_stats(fibpi_path)
        if stats:
            print(f"Index stats for {fibpi_path}:")
            for k, v in stats.items():
                print(f"  {k}: {v}")
        else:
            print(f"No valid index for {fibpi_path}")

    return 0


if __name__ == '__main__':
    exit(main())
