#!/usr/bin/env python3
"""
FibPi SHA Anchor Encoder/Decoder

Encodes data with SHA256 anchors at Fibonacci positions.
Structure: [SHA_0][b64_chunk_0][SHA_1][b64_chunk_1][SHA_2][b64_chunk_2]...

Each anchor is a SHA256 hash (32 bytes hex = 64 chars) of the following chunk.
Chunks are split at cumulative Fibonacci positions and base64 encoded.

Fibonacci positions: 1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233...
Cumulative:          1, 2, 4, 7, 12, 20, 33, 54, 88, 143, 232, 376, 609...

Usage:
    # Encode
    python3 fibpi_sha_anchor.py encode input.bin output.fibpi

    # Decode
    python3 fibpi_sha_anchor.py decode input.fibpi output.bin

    # Verify (check anchors without full decode)
    python3 fibpi_sha_anchor.py verify input.fibpi
"""

import hashlib
import base64
import json
import sys
import os
from pathlib import Path
from datetime import datetime
from typing import List, Tuple, Optional, BinaryIO


# Constants
ANCHOR_PREFIX = b"::SHA::"      # 7 bytes - marks start of anchor
ANCHOR_LEN = 64                  # SHA256 hex length
CHUNK_PREFIX = b"::B64::"       # 7 bytes - marks start of base64 chunk
CHUNK_SUFFIX = b"::END::"       # 7 bytes - marks end of chunk
HEADER_MAGIC = b"FIBPI3D_SHA_v1" # 14 bytes


def generate_fibonacci(max_val: int = 10_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_positions(data_len: int) -> List[int]:
    """Get cumulative Fibonacci positions for chunking."""
    fib = generate_fibonacci(data_len)
    positions = []
    cumulative = 0

    for f in fib:
        cumulative += f
        if cumulative >= data_len:
            positions.append(data_len)
            break
        positions.append(cumulative)

    # If data is longer than our Fibonacci coverage, add final position
    if positions[-1] < data_len:
        positions.append(data_len)

    return positions


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


class FibPiShaEncoder:
    """Encodes data with SHA256 anchors at Fibonacci positions."""

    def __init__(self):
        self.fib_sequence = generate_fibonacci()
        self.chunks_encoded = 0
        self.total_anchors = 0

    def encode(self, data: bytes) -> bytes:
        """
        Encode data with SHA anchors at Fibonacci positions.

        Returns encoded bytes with structure:
        [HEADER][META_JSON][ANCHOR_0][CHUNK_0][ANCHOR_1][CHUNK_1]...
        """
        data_len = len(data)
        positions = get_fibonacci_positions(data_len)

        # Build chunks
        chunks = []
        prev_pos = 0

        for pos in positions:
            chunk = data[prev_pos:pos]
            if chunk:  # Skip empty chunks
                chunks.append(chunk)
            prev_pos = pos

        # Build encoded output
        output = bytearray()

        # Header
        output.extend(HEADER_MAGIC)

        # Metadata (JSON)
        meta = {
            "version": 1,
            "original_size": data_len,
            "num_chunks": len(chunks),
            "fibonacci_positions": positions[:len(chunks)],
            "timestamp": datetime.utcnow().isoformat() + "Z",
            "original_sha256": sha256_hex(data)
        }
        meta_json = json.dumps(meta).encode()
        meta_len = len(meta_json)

        # Write meta length (4 bytes) + meta
        output.extend(meta_len.to_bytes(4, 'big'))
        output.extend(meta_json)

        # Write chunks with anchors
        for i, chunk in enumerate(chunks):
            # SHA256 anchor of the raw chunk
            anchor = sha256_hex(chunk)

            # Base64 encode the chunk
            b64_chunk = base64.b64encode(chunk)

            # Write: [ANCHOR_PREFIX][anchor][CHUNK_PREFIX][b64_data][CHUNK_SUFFIX]
            output.extend(ANCHOR_PREFIX)
            output.extend(anchor.encode())
            output.extend(CHUNK_PREFIX)
            output.extend(b64_chunk)
            output.extend(CHUNK_SUFFIX)

            self.chunks_encoded += 1
            self.total_anchors += 1

        return bytes(output)

    def encode_file(self, input_path: Path, output_path: Path) -> dict:
        """Encode a file and return stats."""
        with open(input_path, 'rb') as f:
            data = f.read()

        encoded = self.encode(data)

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

        original_size = len(data)
        encoded_size = len(encoded)

        return {
            "input": str(input_path),
            "output": str(output_path),
            "original_size": original_size,
            "encoded_size": encoded_size,
            "ratio": round(encoded_size / original_size, 3) if original_size > 0 else 0,
            "chunks": self.chunks_encoded,
            "anchors": self.total_anchors,
            "original_sha256": sha256_hex(data)
        }


class FibPiShaDecoder:
    """Decodes FibPi SHA anchor encoded data."""

    def __init__(self):
        self.chunks_decoded = 0
        self.anchors_verified = 0
        self.anchor_failures = 0

    def decode(self, encoded: bytes, verify: bool = True) -> Tuple[bytes, dict]:
        """
        Decode FibPi SHA anchor encoded data.

        Returns (decoded_data, metadata_dict)
        Raises ValueError if verification fails.
        """
        # Check header
        if not encoded.startswith(HEADER_MAGIC):
            raise ValueError("Invalid FibPi SHA header - not a valid encoded file")

        pos = len(HEADER_MAGIC)

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

        # Read metadata
        meta_json = encoded[pos:pos+meta_len]
        meta = json.loads(meta_json.decode())
        pos += meta_len

        # Decode chunks
        chunks = []

        while pos < len(encoded):
            # Find anchor prefix
            if encoded[pos:pos+len(ANCHOR_PREFIX)] != ANCHOR_PREFIX:
                break
            pos += len(ANCHOR_PREFIX)

            # Read anchor (64 hex chars)
            anchor = encoded[pos:pos+ANCHOR_LEN].decode()
            pos += ANCHOR_LEN

            # Find chunk prefix
            if encoded[pos:pos+len(CHUNK_PREFIX)] != CHUNK_PREFIX:
                raise ValueError(f"Expected chunk prefix at position {pos}")
            pos += len(CHUNK_PREFIX)

            # Find chunk suffix
            suffix_pos = encoded.find(CHUNK_SUFFIX, pos)
            if suffix_pos == -1:
                raise ValueError(f"Missing chunk suffix after position {pos}")

            # Extract base64 chunk
            b64_chunk = encoded[pos:suffix_pos]
            pos = suffix_pos + len(CHUNK_SUFFIX)

            # Decode base64
            try:
                chunk = base64.b64decode(b64_chunk)
            except Exception as e:
                raise ValueError(f"Base64 decode failed: {e}")

            # Verify anchor
            if verify:
                computed_anchor = sha256_hex(chunk)
                if computed_anchor != anchor:
                    self.anchor_failures += 1
                    raise ValueError(
                        f"Anchor verification failed for chunk {len(chunks)}!\n"
                        f"  Expected: {anchor}\n"
                        f"  Computed: {computed_anchor}"
                    )
                self.anchors_verified += 1

            chunks.append(chunk)
            self.chunks_decoded += 1

        # Reassemble
        data = b''.join(chunks)

        # Verify final hash
        if verify and 'original_sha256' in meta:
            final_hash = sha256_hex(data)
            if final_hash != meta['original_sha256']:
                raise ValueError(
                    f"Final SHA256 verification failed!\n"
                    f"  Expected: {meta['original_sha256']}\n"
                    f"  Computed: {final_hash}"
                )

        return data, meta

    def decode_file(self, input_path: Path, output_path: Path, verify: bool = True) -> dict:
        """Decode a file and return stats."""
        with open(input_path, 'rb') as f:
            encoded = f.read()

        data, meta = self.decode(encoded, verify=verify)

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

        return {
            "input": str(input_path),
            "output": str(output_path),
            "encoded_size": len(encoded),
            "decoded_size": len(data),
            "chunks_decoded": self.chunks_decoded,
            "anchors_verified": self.anchors_verified,
            "anchor_failures": self.anchor_failures,
            "metadata": meta
        }

    def verify_file(self, input_path: Path) -> dict:
        """Verify a file without writing output."""
        with open(input_path, 'rb') as f:
            encoded = f.read()

        try:
            data, meta = self.decode(encoded, verify=True)
            return {
                "status": "OK",
                "file": str(input_path),
                "chunks_verified": self.chunks_decoded,
                "anchors_verified": self.anchors_verified,
                "original_size": meta.get('original_size'),
                "original_sha256": meta.get('original_sha256'),
                "computed_sha256": sha256_hex(data),
                "match": sha256_hex(data) == meta.get('original_sha256')
            }
        except ValueError as e:
            return {
                "status": "FAILED",
                "file": str(input_path),
                "error": str(e),
                "chunks_before_failure": self.chunks_decoded,
                "anchor_failures": self.anchor_failures
            }


class FibPiRandomAccess:
    """
    Random access to FibPi SHA anchor encoded files.
    Allows querying/computing without full decompression.
    """

    def __init__(self, filepath: Path):
        self.filepath = Path(filepath)
        self.meta = None
        self.chunk_offsets = []  # (start, end, anchor) for each chunk
        self._index_file()

    def _index_file(self):
        """Build index of chunk positions for random access."""
        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
        chunk_idx = 0
        while pos < len(data):
            if data[pos:pos+len(ANCHOR_PREFIX)] != ANCHOR_PREFIX:
                break

            chunk_start = pos
            pos += len(ANCHOR_PREFIX)

            # Read anchor
            anchor = data[pos:pos+ANCHOR_LEN].decode()
            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((chunk_start, pos, anchor))
            chunk_idx += 1

    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})")

        start, end, expected_anchor = self.chunk_offsets[idx]

        with open(self.filepath, 'rb') as f:
            f.seek(start + len(ANCHOR_PREFIX) + ANCHOR_LEN + len(CHUNK_PREFIX))
            # Read until CHUNK_SUFFIX
            chunk_data = b''
            while True:
                byte = f.read(1)
                if not byte:
                    break
                chunk_data += byte
                if chunk_data.endswith(CHUNK_SUFFIX):
                    chunk_data = chunk_data[:-len(CHUNK_SUFFIX)]
                    break

        # Decode base64
        decoded = base64.b64decode(chunk_data)

        # Verify if requested
        if verify:
            computed = sha256_hex(decoded)
            if computed != expected_anchor:
                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."""
        for idx, (start, end, chunk_anchor) in enumerate(self.chunk_offsets):
            if chunk_anchor == anchor or chunk_anchor.startswith(anchor):
                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 without full decode.
        Returns list of (chunk_index, offset_in_chunk) matches.
        """
        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.
        Only decodes the necessary chunks.
        """
        positions = self.meta.get('fibonacci_positions', [])
        result = bytearray()

        # Find which chunks we need
        prev_pos = 0
        for idx, pos in enumerate(positions):
            chunk_start = prev_pos
            chunk_end = pos

            # Check if this chunk overlaps our range
            if chunk_end > start and chunk_start < end:
                chunk = self.get_chunk_by_index(idx, verify=False)

                # Calculate slice within chunk
                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 [anchor for _, _, anchor in self.chunk_offsets]

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


def encode_with_compression(data: bytes, compression: str = "zlib") -> bytes:
    """
    Encode with optional compression before FibPi SHA encoding.

    compression options:
    - "none": No compression
    - "zlib": Standard zlib compression
    - "cdna": CDNA tensor quantization (for model weights only)
    """
    import zlib

    if compression == "none":
        compressed = data
        comp_flag = b'\x00'
    elif compression == "zlib":
        compressed = zlib.compress(data, level=9)
        comp_flag = b'\x01'
    elif compression == "cdna":
        # CDNA would require tensor data - for now fall back to zlib
        # In production, this would call the CDNA encoder
        compressed = zlib.compress(data, level=9)
        comp_flag = b'\x02'  # Flag for future CDNA
    else:
        raise ValueError(f"Unknown compression: {compression}")

    # Prepend compression flag
    encoder = FibPiShaEncoder()
    return comp_flag + encoder.encode(compressed)


def decode_with_compression(encoded: bytes) -> bytes:
    """Decode with automatic decompression."""
    import zlib

    comp_flag = encoded[0:1]
    encoded_data = encoded[1:]

    decoder = FibPiShaDecoder()
    data, meta = decoder.decode(encoded_data)

    if comp_flag == b'\x00':
        return data
    elif comp_flag == b'\x01':
        return zlib.decompress(data)
    elif comp_flag == b'\x02':
        # CDNA decode - fall back to zlib for now
        return zlib.decompress(data)
    else:
        raise ValueError(f"Unknown compression flag: {comp_flag}")


def main():
    import argparse

    parser = argparse.ArgumentParser(
        description='FibPi SHA Anchor Encoder/Decoder'
    )
    parser.add_argument('command', choices=['encode', 'decode', 'verify'],
                        help='Command to run')
    parser.add_argument('input', help='Input file path')
    parser.add_argument('output', nargs='?', help='Output file path (not needed for verify)')
    parser.add_argument('--no-verify', action='store_true',
                        help='Skip anchor verification during decode')
    parser.add_argument('--compress', choices=['none', 'zlib', 'cdna'], default='none',
                        help='Compression to apply before encoding')

    args = parser.parse_args()

    input_path = Path(args.input)

    if not input_path.exists():
        print(f"Error: Input file not found: {input_path}")
        sys.exit(1)

    if args.command == 'encode':
        if not args.output:
            args.output = str(input_path) + '.fibpi'

        output_path = Path(args.output)

        if args.compress != 'none':
            # Use compression wrapper
            with open(input_path, 'rb') as f:
                data = f.read()
            encoded = encode_with_compression(data, args.compress)
            with open(output_path, 'wb') as f:
                f.write(encoded)

            print(f"Encoded with {args.compress} compression: {output_path}")
            print(f"  Original: {len(data):,} bytes")
            print(f"  Encoded:  {len(encoded):,} bytes")
            print(f"  Ratio:    {len(encoded)/len(data):.3f}x")
        else:
            encoder = FibPiShaEncoder()
            result = encoder.encode_file(input_path, output_path)

            print(f"Encoded: {result['output']}")
            print(f"  Original: {result['original_size']:,} bytes")
            print(f"  Encoded:  {result['encoded_size']:,} bytes")
            print(f"  Ratio:    {result['ratio']}x")
            print(f"  Chunks:   {result['chunks']}")
            print(f"  Anchors:  {result['anchors']}")
            print(f"  SHA256:   {result['original_sha256']}")

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

        output_path = Path(args.output)

        # Check if compressed
        with open(input_path, 'rb') as f:
            first_byte = f.read(1)

        if first_byte in (b'\x00', b'\x01', b'\x02'):
            # Compressed format
            with open(input_path, 'rb') as f:
                encoded = f.read()
            data = decode_with_compression(encoded)
            with open(output_path, 'wb') as f:
                f.write(data)
            print(f"Decoded (with decompression): {output_path}")
            print(f"  Size: {len(data):,} bytes")
        else:
            decoder = FibPiShaDecoder()
            result = decoder.decode_file(input_path, output_path, verify=not args.no_verify)

            print(f"Decoded: {result['output']}")
            print(f"  Encoded:  {result['encoded_size']:,} bytes")
            print(f"  Decoded:  {result['decoded_size']:,} bytes")
            print(f"  Chunks:   {result['chunks_decoded']}")
            print(f"  Verified: {result['anchors_verified']} anchors")

    elif args.command == 'verify':
        decoder = FibPiShaDecoder()
        result = decoder.verify_file(input_path)

        if result['status'] == 'OK':
            print(f"VERIFIED: {result['file']}")
            print(f"  Chunks:   {result['chunks_verified']}")
            print(f"  Anchors:  {result['anchors_verified']}")
            print(f"  SHA256:   {result['original_sha256']}")
            print(f"  Match:    {'YES' if result['match'] else 'NO'}")
        else:
            print(f"FAILED: {result['file']}")
            print(f"  Error:    {result['error']}")
            print(f"  Chunks before failure: {result['chunks_before_failure']}")
            sys.exit(1)


if __name__ == '__main__':
    main()
