#!/usr/bin/env python3
"""
Singular2.py - CCT Empty Space Compression with Checksum Verification
Compresses binary files by detecting and collapsing "empty space" (low-resistance blocks)
while protecting "solid data" (high-resistance blocks).

Usage:
    python singular2.py compress <input.bin> [output.cct2]
    python singular2.py decompress <input.cct2> [output.bin]
"""

import numpy as np
import zlib
import os
import sys
import struct
import hashlib
from typing import Tuple, List, Dict
from pathlib import Path


class CCT_File_Compressor:
    """
    CCT Empty Space Compression Engine
    Detects "empty space" (compressible blocks) and collapses them,
    while protecting "solid data" (incompressible blocks).
    
    File Format (.cct2):
    [4B: magic "CCT2"][4B: block_size][8B: original_size][32B: SHA256]
    [8B: num_blocks][block_data...]
    
    Block Data:
    [1B: type][4B: data_len][data...]
    Type: 0x00 = COLLAPSED (zlib compressed), 0x01 = STATIONARY (raw)
    """
    
    MAGIC = b'CCT2'
    BLOCK_COLLAPSED = 0x00
    BLOCK_STATIONARY = 0x01

    def __init__(self, entropy_threshold: float = 0.27, resistance_threshold: float = 0.5, block_size: int = 4096):
        self.entropy_threshold = entropy_threshold
        self.resistance_threshold = resistance_threshold
        self.block_size = block_size

    def calculate_block_entropy(self, data: bytes) -> float:
        """Calculate Shannon Entropy of a block (0-1 scale)"""
        if len(data) == 0:
            return 0.0
        unique, counts = np.unique(np.frombuffer(data, dtype=np.uint8), return_counts=True)
        probs = counts / len(data)
        entropy = -np.sum(probs * np.log2(probs))
        return entropy / 8.0  # Normalize to 0-1 (max 8 bits per byte)

    def test_compression_resistance(self, data: bytes) -> Tuple[float, bytes]:
        """
        Truth Identification: How much does this block resist compression?
        Low Resistance = Empty Space (Probability)
        High Resistance = Solid Data (Stationary)
        Returns: (resistance_ratio, compressed_data)
        """
        original_size = len(data)
        try:
            compressed = zlib.compress(data, level=9)
            compressed_size = len(compressed)
        except Exception:
            compressed_size = original_size
            compressed = data

        # Resistance = 1 - Compression Ratio
        resistance = compressed_size / original_size
        return resistance, compressed

    def calculate_checksum(self, data: bytes) -> bytes:
        """Calculate SHA256 checksum of data"""
        return hashlib.sha256(data).digest()

    def compress_file(self, input_path: str, output_path: str = None) -> Dict:
        """
        Compress a binary file using CCT Empty Space detection
        """
        if not os.path.exists(input_path):
            raise FileNotFoundError(f"Input file not found: {input_path}")

        if output_path is None:
            output_path = str(Path(input_path).with_suffix('.cct2'))

        # Read input file
        with open(input_path, 'rb') as f:
            file_data = f.read()

        original_size = len(file_data)
        original_checksum = self.calculate_checksum(file_data)

        print(f"--- CCT File Compression: {input_path} ---")
        print(f"Original Size: {original_size} bytes")
        print(f"SHA256: {original_checksum.hex()}")

        total_blocks = (len(file_data) + self.block_size - 1) // self.block_size
        compressed_output = b''
        block_metadata = []

        empty_space_found = 0
        solid_data_found = 0

        for i in range(total_blocks):
            start = i * self.block_size
            end = min((i + 1) * self.block_size, len(file_data))
            block = file_data[start:end]

            if len(block) == 0:
                break

            # 1. Measure Entropy
            block_entropy = self.calculate_block_entropy(block)

            # 2. Test Resistance
            resistance, compressed_block = self.test_compression_resistance(block)

            # 3. CCT Decision: Collapse or Protect?
            if resistance < self.resistance_threshold:
                # COLLAPSE: Use compressed version (Empty Space)
                block_type = self.BLOCK_COLLAPSED
                block_data = compressed_block
                empty_space_found += len(block)
            else:
                # PROTECT: Keep raw (Solid Data)
                block_type = self.BLOCK_STATIONARY
                block_data = block
                solid_data_found += len(block)

            # Store block metadata
            block_metadata.append({
                'type': 'COLLAPSED' if block_type == self.BLOCK_COLLAPSED else 'STATIONARY',
                'entropy': float(block_entropy),
                'resistance': float(resistance)
            })

            # Write block to output: [1B: type][4B: data_len][data...]
            compressed_output += struct.pack('<B', block_type)
            compressed_output += struct.pack('<I', len(block_data))
            compressed_output += block_data

        # Calculate final size with header
        header_size = 4 + 4 + 8 + 32 + 8  # magic + block_size + original_size + checksum + num_blocks
        final_size = header_size + len(compressed_output)

        # Write compressed file with header
        with open(output_path, 'wb') as f:
            # [4B: magic "CCT2"]
            f.write(self.MAGIC)
            # [4B: block_size]
            f.write(struct.pack('<I', self.block_size))
            # [8B: original_size]
            f.write(struct.pack('<Q', original_size))
            # [32B: SHA256 checksum]
            f.write(original_checksum)
            # [8B: num_blocks]
            f.write(struct.pack('<Q', len(block_metadata)))
            # [block_data...]
            f.write(compressed_output)

        compression_ratio = final_size / original_size if original_size > 0 else 0
        space_saved = (1 - compression_ratio) * 100

        print(f"Final Size: {final_size} bytes")
        print(f"Compression Ratio: {compression_ratio:.4f}")
        print(f"Space Saved: {space_saved:.1f}%")
        print(f"Empty Space Collapsed: {empty_space_found} bytes ({empty_space_found/original_size*100:.1f}%)")
        print(f"Solid Data Protected: {solid_data_found} bytes ({solid_data_found/original_size*100:.1f}%)")
        print(f"Total Blocks: {len(block_metadata)}")
        print(f"Output: {output_path}")
        print("-" * 70)

        # Print block summary
        collapsed_count = sum(1 for m in block_metadata if m['type'] == 'COLLAPSED')
        stationary_count = sum(1 for m in block_metadata if m['type'] == 'STATIONARY')
        print(f"Block Summary: {collapsed_count} collapsed, {stationary_count} stationary")

        return {
            'original_size': original_size,
            'compressed_size': final_size,
            'compression_ratio': compression_ratio,
            'space_saved_percent': space_saved,
            'original_checksum': original_checksum,
            'blocks_processed': len(block_metadata),
            'blocks_collapsed': collapsed_count,
            'blocks_stationary': stationary_count
        }

    def decompress_file(self, input_path: str, output_path: str = None) -> Dict:
        """
        Decompress a .cct2 file back to original binary data with checksum verification
        """
        if not os.path.exists(input_path):
            raise FileNotFoundError(f"Input file not found: {input_path}")

        if output_path is None:
            # Remove .cct2 extension and restore original
            stem = Path(input_path).stem
            output_path = str(Path(input_path).parent / stem)

        with open(input_path, 'rb') as f:
            # [4B: magic "CCT2"]
            magic = f.read(4)
            if magic != self.MAGIC:
                raise ValueError(f"Invalid file format. Expected magic {self.MAGIC}, got {magic}")

            # [4B: block_size]
            block_size = struct.unpack('<I', f.read(4))[0]

            # [8B: original_size]
            original_size = struct.unpack('<Q', f.read(8))[0]

            # [32B: SHA256 checksum]
            expected_checksum = f.read(32)

            # [8B: num_blocks]
            num_blocks = struct.unpack('<Q', f.read(8))[0]

            # Read and reconstruct blocks
            decompressed_output = b''
            blocks_collapsed = 0
            blocks_stationary = 0

            for i in range(num_blocks):
                # [1B: type]
                block_type = struct.unpack('<B', f.read(1))[0]

                # [4B: data_len]
                data_len = struct.unpack('<I', f.read(4))[0]

                # [data...]
                block_data = f.read(data_len)

                if block_type == self.BLOCK_COLLAPSED:
                    # Decompress collapsed block
                    try:
                        decompressed_block = zlib.decompress(block_data)
                        blocks_collapsed += 1
                    except Exception as e:
                        raise ValueError(f"Failed to decompress block {i}: {str(e)}")
                elif block_type == self.BLOCK_STATIONARY:
                    # Keep raw
                    decompressed_block = block_data
                    blocks_stationary += 1
                else:
                    raise ValueError(f"Unknown block type: {block_type} at block {i}")

                decompressed_output += decompressed_block

        # Verify checksum
        actual_checksum = self.calculate_checksum(decompressed_output)
        checksum_match = actual_checksum == expected_checksum

        # Verify size
        size_match = len(decompressed_output) == original_size

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

        print(f"--- CCT File Decompression: {input_path} ---")
        print(f"Compressed Size: {os.path.getsize(input_path)} bytes")
        print(f"Decompressed Size: {len(decompressed_output)} bytes")
        print(f"Expected Size: {original_size} bytes")
        print(f"Blocks: {blocks_collapsed} collapsed, {blocks_stationary} stationary")
        print(f"Output: {output_path}")
        print("-" * 70)

        # Verification results
        print("Verification:")
        if size_match:
            print(f"  [✓] Size verification: PASSED ({len(decompressed_output)} bytes)")
        else:
            print(f"  [✗] Size verification: FAILED (Expected {original_size}, got {len(decompressed_output)})")

        if checksum_match:
            print(f"  [✓] Checksum verification: PASSED")
            print(f"      SHA256: {actual_checksum.hex()}")
        else:
            print(f"  [✗] Checksum verification: FAILED")
            print(f"      Expected: {expected_checksum.hex()}")
            print(f"      Actual:   {actual_checksum.hex()}")

        print("=" * 70)

        if not checksum_match:
            raise ValueError("Checksum verification failed! File may be corrupted.")

        return {
            'decompressed_size': len(decompressed_output),
            'expected_size': original_size,
            'checksum_match': checksum_match,
            'size_match': size_match,
            'actual_checksum': actual_checksum,
            'expected_checksum': expected_checksum,
            'blocks_collapsed': blocks_collapsed,
            'blocks_stationary': blocks_stationary,
            'output_path': output_path
        }


def compress_file_cli(input_path: str, output_path: str = None):
    """CLI wrapper for compression"""
    compressor = CCT_File_Compressor()
    try:
        result = compressor.compress_file(input_path, output_path)
        print(f"\n[✓] Compression successful!")
        return result
    except Exception as e:
        print(f"\n[✗] Compression failed: {str(e)}")
        raise


def decompress_file_cli(input_path: str, output_path: str = None):
    """CLI wrapper for decompression"""
    compressor = CCT_File_Compressor()
    try:
        result = compressor.decompress_file(input_path, output_path)
        print(f"\n[✓] Decompression successful!")
        return result
    except Exception as e:
        print(f"\n[✗] Decompression failed: {str(e)}")
        raise


def main():
    """CLI entry point"""
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)

    command = sys.argv[1].lower()

    if command == 'compress':
        if len(sys.argv) < 3:
            print("Usage: python singular2.py compress <input.bin> [output.cct2]")
            sys.exit(1)

        input_file = sys.argv[2]
        output_file = sys.argv[3] if len(sys.argv) > 3 else None

        compress_file_cli(input_file, output_file)

    elif command == 'decompress':
        if len(sys.argv) < 3:
            print("Usage: python singular2.py decompress <input.cct2> [output.bin]")
            sys.exit(1)

        input_file = sys.argv[2]
        output_file = sys.argv[3] if len(sys.argv) > 3 else None

        decompress_file_cli(input_file, output_file)

    else:
        print(f"Unknown command: {command}")
        print("Available commands: compress, decompress")
        sys.exit(1)


if __name__ == '__main__':
    main()