#!/usr/bin/env python3
"""
FibPi Parallel Decoder - Multi-threaded chunk decoding for FibPi SHA Anchor files

Uses ThreadPoolExecutor and memory-mapped I/O for fast parallel chunk decoding.
Target: 4x speedup on 4-core system for files with 20+ chunks.

Usage:
    from fibpi_parallel import ParallelChunkDecoder

    # Decode file with 4 workers
    decoder = ParallelChunkDecoder("data.fibpi", max_workers=4)
    data = decoder.decode_all_parallel()

    # Benchmark parallel vs sequential
    python3 fibpi_parallel.py benchmark data.fibpi
"""

import os
import sys
import mmap
import base64
import hashlib
import json
import time
from pathlib import Path
from typing import List, Tuple, Optional, Dict, Any
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
import threading

# Import base FibPi components - try relative, then absolute, then fallback
FibPiIndexManager = None

try:
    # Try relative import (package context)
    from .fibpi_sha_anchor import (
        HEADER_MAGIC, ANCHOR_PREFIX, ANCHOR_LEN, CHUNK_PREFIX, CHUNK_SUFFIX,
        FibPiShaDecoder, sha256_hex
    )
    from .fibpi_index import FibPiIndexManager, ChunkInfo
except ImportError:
    try:
        # Try absolute import (direct execution)
        from fibpi_sha_anchor import (
            HEADER_MAGIC, ANCHOR_PREFIX, ANCHOR_LEN, CHUNK_PREFIX, CHUNK_SUFFIX,
            FibPiShaDecoder, sha256_hex
        )
        from fibpi_index import FibPiIndexManager, ChunkInfo
    except ImportError:
        # Fallback constants 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()

        @dataclass
        class ChunkInfo:
            offset: int
            length: int
            anchor: bytes

            @property
            def anchor_hex(self) -> str:
                return self.anchor.hex()


class MMapReader:
    """Memory-mapped file reader for zero-copy access."""

    def __init__(self, filepath: Path):
        self.filepath = Path(filepath)
        self.file = None
        self.mm = None
        self._lock = threading.Lock()

    def __enter__(self):
        self.file = open(self.filepath, 'rb')
        self.mm = mmap.mmap(self.file.fileno(), 0, access=mmap.ACCESS_READ)
        return self

    def __exit__(self, *args):
        if self.mm:
            self.mm.close()
        if self.file:
            self.file.close()

    def read_range(self, start: int, length: int) -> bytes:
        """Thread-safe read of a byte range."""
        with self._lock:
            return self.mm[start:start + length]


class ParallelChunkDecoder:
    """
    Decode multiple FibPi chunks in parallel using thread pool.

    Uses memory-mapped I/O for efficient concurrent reads without
    file descriptor conflicts.
    """

    def __init__(self, fibpi_path: Path, max_workers: int = 4, use_index: bool = True):
        """
        Initialize parallel decoder.

        Args:
            fibpi_path: Path to .fibpi encoded file
            max_workers: Maximum number of parallel decode threads
            use_index: Whether to use sidecar index for fast chunk lookup
        """
        self.filepath = Path(fibpi_path)
        self.max_workers = max_workers
        self.chunk_infos: List[ChunkInfo] = []
        self.meta: Dict[str, Any] = {}

        # Try to load index
        if use_index:
            try:
                mgr = FibPiIndexManager()
                cached = mgr.load_index(self.filepath)
                if cached:
                    self.chunk_infos = cached
                    self._load_metadata()
                else:
                    self._index_file()
                    mgr.create_index(self.filepath)
            except Exception:
                self._index_file()
        else:
            self._index_file()

        # Stats
        self.stats = {
            'chunks_decoded': 0,
            'bytes_decoded': 0,
            'verification_failures': 0,
            'parallel_time_ms': 0,
            'sequential_time_ms': 0
        }

    def _load_metadata(self):
        """Load metadata from FibPi file."""
        with open(self.filepath, 'rb') as f:
            magic = f.read(len(HEADER_MAGIC))
            if magic != HEADER_MAGIC:
                raise ValueError("Invalid FibPi file")

            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 chunk index by scanning file."""
        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 chunks
        while pos < len(data):
            if data[pos:pos+len(ANCHOR_PREFIX)] != ANCHOR_PREFIX:
                break

            chunk_start = pos
            pos += len(ANCHOR_PREFIX)

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

            suffix_pos = data.find(CHUNK_SUFFIX, pos)
            if suffix_pos == -1:
                break
            pos = suffix_pos + len(CHUNK_SUFFIX)

            self.chunk_infos.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_infos)

    def _decode_single_chunk(self, mm: mmap.mmap, idx: int, verify: bool) -> Tuple[int, bytes]:
        """
        Decode a single chunk from memory-mapped file.
        Thread-safe via mmap's internal synchronization.

        Returns (chunk_index, decoded_bytes)
        """
        info = self.chunk_infos[idx]

        # Calculate positions within chunk
        b64_start = info.offset + len(ANCHOR_PREFIX) + ANCHOR_LEN + len(CHUNK_PREFIX)
        b64_len = info.length - len(ANCHOR_PREFIX) - ANCHOR_LEN - len(CHUNK_PREFIX) - len(CHUNK_SUFFIX)

        # Read base64 data
        b64_data = mm[b64_start:b64_start + b64_len]

        # Decode
        decoded = base64.b64decode(b64_data)

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

        return (idx, decoded)

    def decode_all_parallel(self, verify: bool = True) -> bytes:
        """
        Decode all chunks in parallel, assemble in order.

        Args:
            verify: Whether to verify SHA256 anchors

        Returns:
            Original file bytes
        """
        if not self.chunk_infos:
            return b''

        start_time = time.perf_counter()
        results: List[Optional[bytes]] = [None] * len(self.chunk_infos)
        errors = []

        # Use memory-mapped file for thread-safe concurrent reads
        with open(self.filepath, 'rb') as f:
            mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)

            try:
                with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                    # Submit all chunk decode tasks
                    futures = {
                        executor.submit(self._decode_single_chunk, mm, i, verify): i
                        for i in range(len(self.chunk_infos))
                    }

                    # Collect results
                    for future in as_completed(futures):
                        idx = futures[future]
                        try:
                            chunk_idx, data = future.result()
                            results[chunk_idx] = data
                            self.stats['chunks_decoded'] += 1
                            self.stats['bytes_decoded'] += len(data)
                        except Exception as e:
                            errors.append((idx, str(e)))
                            self.stats['verification_failures'] += 1

            finally:
                mm.close()

        self.stats['parallel_time_ms'] = (time.perf_counter() - start_time) * 1000

        if errors:
            error_msg = "; ".join(f"chunk {idx}: {msg}" for idx, msg in errors)
            raise ValueError(f"Parallel decode failed: {error_msg}")

        # Verify all chunks decoded
        if None in results:
            missing = [i for i, r in enumerate(results) if r is None]
            raise ValueError(f"Missing chunks: {missing}")

        return b''.join(results)

    def decode_sequential(self, verify: bool = True) -> bytes:
        """
        Decode all chunks sequentially (for comparison).
        """
        if not self.chunk_infos:
            return b''

        start_time = time.perf_counter()
        results = []

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

        for idx, info in enumerate(self.chunk_infos):
            b64_start = info.offset + len(ANCHOR_PREFIX) + ANCHOR_LEN + len(CHUNK_PREFIX)
            b64_len = info.length - len(ANCHOR_PREFIX) - ANCHOR_LEN - len(CHUNK_PREFIX) - len(CHUNK_SUFFIX)

            b64_data = data[b64_start:b64_start + b64_len]
            decoded = base64.b64decode(b64_data)

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

            results.append(decoded)

        self.stats['sequential_time_ms'] = (time.perf_counter() - start_time) * 1000
        return b''.join(results)

    def decode_chunks_parallel(self, chunk_indices: List[int], verify: bool = True) -> Dict[int, bytes]:
        """
        Decode specific chunks in parallel.

        Args:
            chunk_indices: List of chunk indices to decode
            verify: Whether to verify anchors

        Returns:
            Dict mapping chunk index to decoded bytes
        """
        results = {}

        with open(self.filepath, 'rb') as f:
            mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)

            try:
                with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                    futures = {
                        executor.submit(self._decode_single_chunk, mm, i, verify): i
                        for i in chunk_indices if 0 <= i < len(self.chunk_infos)
                    }

                    for future in as_completed(futures):
                        idx = futures[future]
                        try:
                            chunk_idx, data = future.result()
                            results[chunk_idx] = data
                        except Exception as e:
                            results[idx] = None

            finally:
                mm.close()

        return results

    def benchmark(self, iterations: int = 3) -> Dict[str, Any]:
        """
        Benchmark parallel vs sequential decode.

        Returns timing comparison.
        """
        parallel_times = []
        sequential_times = []

        for _ in range(iterations):
            # Parallel
            t0 = time.perf_counter()
            _ = self.decode_all_parallel(verify=False)
            parallel_times.append(time.perf_counter() - t0)

            # Sequential
            t0 = time.perf_counter()
            _ = self.decode_sequential(verify=False)
            sequential_times.append(time.perf_counter() - t0)

        avg_parallel = sum(parallel_times) / len(parallel_times)
        avg_sequential = sum(sequential_times) / len(sequential_times)

        return {
            'parallel_avg_ms': avg_parallel * 1000,
            'sequential_avg_ms': avg_sequential * 1000,
            'speedup': avg_sequential / avg_parallel if avg_parallel > 0 else 0,
            'workers': self.max_workers,
            'chunks': len(self.chunk_infos),
            'iterations': iterations
        }


class BufferedChunkReader:
    """
    Buffered reader for streaming large FibPi files.

    Decodes chunks on-demand with configurable buffer size.
    """

    def __init__(self, filepath: Path, buffer_chunks: int = 4):
        """
        Initialize buffered reader.

        Args:
            filepath: Path to .fibpi file
            buffer_chunks: Number of chunks to buffer ahead
        """
        self.filepath = Path(filepath)
        self.buffer_chunks = buffer_chunks
        self.decoder = ParallelChunkDecoder(filepath, max_workers=buffer_chunks)

        self._buffer: Dict[int, bytes] = {}
        self._current_idx = 0
        self._lock = threading.Lock()

    def read_chunk(self, idx: int) -> bytes:
        """Read a single chunk, using buffer if available."""
        with self._lock:
            if idx in self._buffer:
                return self._buffer[idx]

        # Prefetch buffer_chunks ahead
        prefetch_indices = list(range(idx, min(idx + self.buffer_chunks, self.decoder.get_chunk_count())))
        chunks = self.decoder.decode_chunks_parallel(prefetch_indices, verify=True)

        with self._lock:
            self._buffer.update({k: v for k, v in chunks.items() if v is not None})
            self._evict_old(idx)
            return self._buffer.get(idx, b'')

    def _evict_old(self, current_idx: int):
        """Evict chunks that are too far behind current position."""
        to_evict = [k for k in self._buffer if k < current_idx - 2]
        for k in to_evict:
            del self._buffer[k]

    def read_sequential(self) -> bytes:
        """Read all chunks in order, using buffered prefetch."""
        results = []
        for idx in range(self.decoder.get_chunk_count()):
            results.append(self.read_chunk(idx))
        return b''.join(results)


def decode_file_parallel(input_path: Path, output_path: Path,
                         workers: int = 4, verify: bool = True) -> Dict[str, Any]:
    """
    Convenience function to decode a FibPi file with parallel processing.

    Args:
        input_path: Path to .fibpi file
        output_path: Path for decoded output
        workers: Number of parallel workers
        verify: Whether to verify anchors

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

    decoder = ParallelChunkDecoder(input_path, max_workers=workers)

    start_time = time.perf_counter()
    data = decoder.decode_all_parallel(verify=verify)
    decode_time = time.perf_counter() - start_time

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

    return {
        'input': str(input_path),
        'output': str(output_path),
        'chunks': decoder.get_chunk_count(),
        'workers': workers,
        'decoded_size': len(data),
        'decode_time_ms': decode_time * 1000,
        'throughput_mbps': (len(data) / (1024 * 1024)) / decode_time if decode_time > 0 else 0,
        'verified': verify
    }


def main():
    import argparse

    parser = argparse.ArgumentParser(description='FibPi Parallel Decoder')
    parser.add_argument('command', choices=['decode', 'benchmark'],
                        help='Command to run')
    parser.add_argument('input', help='Input .fibpi file')
    parser.add_argument('output', nargs='?', help='Output file (for decode)')
    parser.add_argument('-w', '--workers', type=int, default=4,
                        help='Number of parallel workers (default: 4)')
    parser.add_argument('--no-verify', action='store_true',
                        help='Skip anchor verification')
    parser.add_argument('-i', '--iterations', type=int, default=3,
                        help='Benchmark iterations (default: 3)')

    args = parser.parse_args()
    input_path = Path(args.input)

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

    if args.command == 'decode':
        if not args.output:
            if str(input_path).endswith('.fibpi'):
                args.output = str(input_path)[:-6]
            else:
                args.output = str(input_path) + '.decoded'

        output_path = Path(args.output)

        result = decode_file_parallel(
            input_path, output_path,
            workers=args.workers,
            verify=not args.no_verify
        )

        print(f"Parallel decode complete: {result['output']}")
        print(f"  Workers:     {result['workers']}")
        print(f"  Chunks:      {result['chunks']}")
        print(f"  Size:        {result['decoded_size']:,} bytes")
        print(f"  Time:        {result['decode_time_ms']:.1f} ms")
        print(f"  Throughput:  {result['throughput_mbps']:.1f} MB/s")
        print(f"  Verified:    {result['verified']}")

    elif args.command == 'benchmark':
        decoder = ParallelChunkDecoder(input_path, max_workers=args.workers)
        results = decoder.benchmark(iterations=args.iterations)

        print(f"Benchmark results for {input_path}:")
        print(f"  Workers:      {results['workers']}")
        print(f"  Chunks:       {results['chunks']}")
        print(f"  Iterations:   {results['iterations']}")
        print(f"  Parallel:     {results['parallel_avg_ms']:.2f} ms")
        print(f"  Sequential:   {results['sequential_avg_ms']:.2f} ms")
        print(f"  Speedup:      {results['speedup']:.2f}x")

    return 0


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