#!/usr/bin/env python3
"""
Singular.py - CCT Self-Learning Binary Compression Engine
Combines compression and decompression in a single unified interface.

Usage:
    python singular.py compress <input.bin> [output.cct]
    python singular.py decompress <input.cct> [output.bin]
"""

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


class CCT_Compression_Engine:
    """
    Self-Learning Compression using CCT-ULP Framework
    Combines: File 1 (CCT), File 2 (SuperByte), File 3 (ODE-CCT),
              File 5 (16-Element), File 6 (Constants)
    """

    def __init__(self, entropy_threshold: float = 0.27):
        self.entropy_threshold = entropy_threshold
        self.entropy_history = []
        self.algorithm_history = []

        # E01: Stationary Constants (Information Laws)
        self.stationary_constants = {
            'shannon_limit': None,
            'kolmogorov_bound': None,
            'causality': True  # Decompression must reverse compression
        }

        # E02: Probability Variables (Algorithm Parameters)
        self.probability_variables = {
            'algorithm': 'auto',  # Will be selected
            'compression_level': 6,
            'block_size': 4096,
            'use_transform': False,
            'dictionary_size': 32768
        }

        # E07/E08: Available Algorithms (Transform Basis)
        self.algorithms = {
            'zlib': lambda d, l: zlib.compress(d, l),
            'bz2': lambda d, l: bz2.compress(d, l),
            'lzma': lambda d, l: lzma.compress(d, preset=l),
            'zstd': lambda d, l: self._zstd_compress(d, l),
            'none': lambda d, l: d  # Baseline
        }

        # Decompression mapping
        self.decompressors = {
            'zlib': zlib.decompress,
            'bz2': bz2.decompress,
            'lzma': lzma.decompress,
            'zstd': self._zstd_decompress,
            'none': lambda d: d
        }

        # 16-Element Compression Matrix
        self.elements = self._initialize_compression_elements()

    def _initialize_compression_elements(self) -> Dict[str, float]:
        """Initialize 16-Element Compression Matrix"""
        return {
            'E01_Entropy_Bound': 0.0,
            'E02_Pattern_Density': 0.0,
            'E03_Redundancy_Rate': 0.0,
            'E04_Algorithm_Select': 0.0,
            'E05_Block_Optimal': 0.0,
            'E06_Dictionary_Size': 0.0,
            'E07_Transform_Type': 0.0,
            'E08_Prediction_Error': 0.0,
            'E09_Entropy_Current': 0.0,
            'E10_Collapse_Rate': 0.0,
            'E11_Memory_Cost': 0.0,
            'E12_Compute_Cost': 0.0,
            'E13_Recovery_Fidelity': 0.0,
            'E14_Context_Window': 0.0,
            'E15_Novelty_Score': 0.0,
            'E16_Compression_Stability': 0.0
        }

    def calculate_compression_entropy(self, data: bytes, params: Dict) -> float:
        """
        E09: Entropy Metric
        Measures compressibility as normalized entropy
        """
        # Get compressed size
        algo = params.get('algorithm', 'zlib')
        level = params.get('compression_level', 6)

        if algo == 'auto':
            algo = 'zlib'  # Default for measurement

        try:
            compressed = self.algorithms[algo](data, level)
            compression_ratio = len(compressed) / len(data)
        except:
            compression_ratio = 1.0

        # Normalize to [0, 1] entropy scale
        # 0 = perfectly compressed, 1 = incompressible
        entropy = np.tanh(compression_ratio)

        self.elements['E09_Entropy_Current'] = entropy
        return entropy

    def test_parameter_resistance(self, data: bytes, param_name: str) -> float:
        """
        Truth Identification: Test what resists bending
        High resistance = Stationary Constant (Truth)
        """
        base_params = {**self.stationary_constants, **self.probability_variables}
        base_entropy = self.calculate_compression_entropy(data, base_params)

        # Perturb parameter
        test_params = base_params.copy()
        if param_name == 'compression_level':
            test_params[param_name] = min(9, test_params[param_name] + 2)
        elif param_name == 'block_size':
            test_params[param_name] = test_params[param_name] * 2
        elif param_name == 'algorithm':
            test_params[param_name] = 'bz2'  # Switch algorithm

        test_entropy = self.calculate_compression_entropy(data, test_params)

        # Resistance = entropy sensitivity
        resistance = abs(test_entropy - base_entropy) / 0.1
        return resistance

    def learn_and_compress(self, data: bytes, max_iterations: int = 10) -> Tuple[Dict, bytes]:
        """
        CCT-ULP Self-Learning Compression Loop
        """
        print("="*70)
        print("CCT SELF-LEARNING COMPRESSION ENGINE")
        print("="*70)

        # Calculate initial entropy
        current_entropy = self.calculate_compression_entropy(data, self.probability_variables)
        self.entropy_history.append(current_entropy)

        print(f"Initial Entropy: {current_entropy:.4f} (Target < {self.entropy_threshold})")
        print(f"Data Size: {len(data)} bytes")
        print("-"*70)

        # CCT Learning Loop
        for iteration in range(max_iterations):
            print(f"\n--- CCT Cycle {iteration + 1} ---")

            # Step 1: Truth Identification
            for param in list(self.probability_variables.keys()):
                if param == 'algorithm':
                    continue  # Skip algorithm selection for resistance test

                resistance = self.test_parameter_resistance(data, param)

                if resistance > 0.5:
                    print(f"  [!] '{param}' has HIGH resistance → Stationary")
                else:
                    print(f"  [~] '{param}' is bendable → Probability")

            # Step 2: Check Stability
            if current_entropy < self.entropy_threshold:
                print(f"\n[✓] System Stable (Entropy < 0.27)")
                
                # If algorithm is still 'auto', select the best one
                if self.probability_variables['algorithm'] == 'auto':
                    print(f"\n[⚠] Algorithm not yet selected. Testing algorithms...")
                    best_algo = None
                    best_entropy = 1.0

                    for algo_name in self.algorithms.keys():
                        if algo_name == 'none':
                            continue

                        test_params = {**self.probability_variables, 'algorithm': algo_name}
                        entropy = self.calculate_compression_entropy(data, test_params)

                        print(f"  {algo_name}: Entropy = {entropy:.4f}")

                        if entropy < best_entropy:
                            best_entropy = entropy
                            best_algo = algo_name

                    if best_algo:
                        self.probability_variables['algorithm'] = best_algo
                        current_entropy = best_entropy
                        self.algorithm_history.append(best_algo)
                        print(f"\n[✓] Selected: {best_algo} (Entropy: {current_entropy:.4f})")
                
                break

            # Step 3: Variable Bending (Algorithm Selection)
            print(f"\n[⚠] Entropy High. Testing algorithms...")

            best_algo = None
            best_entropy = 1.0

            for algo_name in self.algorithms.keys():
                if algo_name == 'none':
                    continue

                test_params = {**self.probability_variables, 'algorithm': algo_name}
                entropy = self.calculate_compression_entropy(data, test_params)

                print(f"  {algo_name}: Entropy = {entropy:.4f}")

                if entropy < best_entropy:
                    best_entropy = entropy
                    best_algo = algo_name

            # Update to best algorithm
            if best_algo:
                self.probability_variables['algorithm'] = best_algo
                current_entropy = best_entropy
                self.algorithm_history.append(best_algo)
                print(f"\n[✓] Selected: {best_algo} (Entropy: {current_entropy:.4f})")

            self.entropy_history.append(current_entropy)

            # Step 4: Singularity Check
            if iteration == max_iterations - 1:
                print(f"\n[!] Max iterations reached. Partial collapse.")

        # E16: Conditional Collapse (Final Compression)
        final_params = {**self.stationary_constants, **self.probability_variables}
        compressed = self.algorithms[final_params['algorithm']](
            data,
            final_params['compression_level']
        )

        # Calculate final metrics
        compression_ratio = len(compressed) / len(data)
        space_saved = (1 - compression_ratio) * 100

        result = {
            'status': 'COLLAPSED' if current_entropy < self.entropy_threshold else 'PARTIAL',
            'algorithm': final_params['algorithm'],
            'original_size': len(data),
            'compressed_size': len(compressed),
            'compression_ratio': compression_ratio,
            'space_saved_percent': space_saved,
            'final_entropy': current_entropy,
            'iterations': iteration + 1,
            'elements': self.elements
        }

        print("\n" + "="*70)
        print("FINAL COMPRESSION METRICS:")
        print(f"  Status: {result['status']}")
        print(f"  Algorithm: {result['algorithm']}")
        print(f"  Original: {result['original_size']} bytes")
        print(f"  Compressed: {result['compressed_size']} bytes")
        print(f"  Ratio: {result['compression_ratio']:.4f}")
        print(f"  Space Saved: {result['space_saved_percent']:.1f}%")
        print(f"  Final Entropy: {result['final_entropy']:.4f}")
        print("="*70)

        return result, compressed

    def decompress(self, data: bytes, algorithm: str) -> bytes:
        """
        Decompress data using specified algorithm
        """
        try:
            decompressor = self.decompressors[algorithm]
            return decompressor(data)
        except Exception as e:
            raise ValueError(f"Decompression failed for '{algorithm}': {str(e)}")

    def _zstd_compress(self, data: bytes, level: int) -> bytes:
        """Zstandard compression (requires zstd library)"""
        try:
            import zstd
            return zstd.compress(data, level)
        except ImportError:
            return zlib.compress(data, level)  # Fallback

    def _zstd_decompress(self, data: bytes) -> bytes:
        """Zstandard decompression (requires zstd library)"""
        try:
            import zstd
            return zstd.decompress(data)
        except ImportError:
            raise ImportError("zstd library not available. Install with: pip install zstd")


class CCT_FileHandler:
    """
    Handles file I/O with custom .cct format
    Format: [4 bytes: algo_name_len][algo_name][8 bytes: original_size][compressed_data]
    """

    @staticmethod
    def save_compressed(compressed_data: bytes, algorithm: str, original_size: int, output_path: str):
        """
        Save compressed data with metadata for decompression
        """
        algo_bytes = algorithm.encode('utf-8')
        algo_len = len(algo_bytes)

        with open(output_path, 'wb') as f:
            # Write algorithm name length (4 bytes, little-endian)
            f.write(struct.pack('<I', algo_len))
            # Write algorithm name
            f.write(algo_bytes)
            # Write original size (8 bytes, little-endian)
            f.write(struct.pack('<Q', original_size))
            # Write compressed data
            f.write(compressed_data)

        print(f"\n[✓] Saved compressed file: {output_path}")

    @staticmethod
    def load_compressed(input_path: str) -> Tuple[bytes, str, int]:
        """
        Load compressed file and extract metadata
        Returns: (compressed_data, algorithm, original_size)
        """
        with open(input_path, 'rb') as f:
            # Read algorithm name length
            algo_len_data = f.read(4)
            if len(algo_len_data) < 4:
                raise ValueError("Invalid compressed file format")
            
            algo_len = struct.unpack('<I', algo_len_data)[0]
            
            # Read algorithm name
            algo_name = f.read(algo_len).decode('utf-8')
            
            # Read original size
            original_size_data = f.read(8)
            if len(original_size_data) < 8:
                raise ValueError("Invalid compressed file format")
            
            original_size = struct.unpack('<Q', original_size_data)[0]
            
            # Read compressed data
            compressed_data = f.read()

        print(f"\n[✓] Loaded compressed file: {input_path}")
        print(f"  Algorithm: {algo_name}")
        print(f"  Original size: {original_size} bytes")
        print(f"  Compressed size: {len(compressed_data)} bytes")

        return compressed_data, algo_name, original_size


def compress_file(input_path: str, output_path: Optional[str] = None, verbose: bool = True):
    """
    Compress a binary file using CCT self-learning compression
    """
    if not os.path.exists(input_path):
        print(f"[✗] File not found: {input_path}")
        return

    # Generate output path if not provided
    if output_path is None:
        output_path = str(Path(input_path).with_suffix('.cct'))

    # Read input file
    print(f"\nReading file: {input_path}")
    with open(input_path, 'rb') as f:
        data = f.read()

    print(f"File size: {len(data)} bytes")

    # Initialize compression engine
    engine = CCT_Compression_Engine()

    # Compress with self-learning
    result, compressed = engine.learn_and_compress(data)

    # Save compressed file with metadata
    CCT_FileHandler.save_compressed(
        compressed_data=compressed,
        algorithm=result['algorithm'],
        original_size=result['original_size'],
        output_path=output_path
    )

    return result


def decompress_file(input_path: str, output_path: Optional[str] = None):
    """
    Decompress a .cct file back to original binary data
    """
    if not os.path.exists(input_path):
        print(f"[✗] File not found: {input_path}")
        return

    # Generate output path if not provided
    if output_path is None:
        # Remove .cct extension and restore original
        output_path = str(Path(input_path).with_suffix(''))

    # Load compressed file
    compressed_data, algorithm, original_size = CCT_FileHandler.load_compressed(input_path)

    # Initialize engine for decompression
    engine = CCT_Compression_Engine()

    # Decompress
    print(f"\nDecompressing with algorithm: {algorithm}")
    decompressed = engine.decompress(compressed_data, algorithm)

    # Verify size
    if len(decompressed) != original_size:
        print(f"[!] Warning: Decompressed size ({len(decompressed)}) != Expected ({original_size})")

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

    print(f"\n[✓] Decompressed file saved: {output_path}")
    print(f"  Size: {len(decompressed)} bytes")

    # Verify integrity
    if len(decompressed) == original_size:
        print(f"[✓] Size verification: PASSED")
    else:
        print(f"[✗] Size verification: FAILED")

    return decompressed


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 singular.py compress <input.bin> [output.cct]")
            sys.exit(1)
        
        input_file = sys.argv[2]
        output_file = sys.argv[3] if len(sys.argv) > 3 else None
        
        compress_file(input_file, output_file)

    elif command == 'decompress':
        if len(sys.argv) < 3:
            print("Usage: python singular.py decompress <input.cct> [output.bin]")
            sys.exit(1)
        
        input_file = sys.argv[2]
        output_file = sys.argv[3] if len(sys.argv) > 3 else None
        
        decompress_file(input_file, output_file)

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


if __name__ == '__main__':
    main()