import numpy as np
import hashlib
import zlib
import bz2
import lzma
import argparse
import tarfile
import io
import os
from pathlib import Path
from typing import List, Dict, Tuple
from dataclasses import dataclass, field

# ============================================================
# DATA STRUCTURES
# ============================================================
@dataclass
class CompressionAlgorithm:
    """Represents one of the 100+ potential algorithms"""
    name: str
    id: int
    category: str  # 'dictionary', 'statistical', 'transform', 'semantic'
    cost_weight: float  # Compute energy cost
    active: bool = True
    success_history: List[float] = field(default_factory=list)
    weight: float = 0.5  # SuperBoolean probability weight

@dataclass
class ChunkResult:
    """Result of compressing a single chunk"""
    chunk_id: int
    algo_id: int
    original_size: int
    compressed_size: int
    entropy_before: float
    entropy_after: float
    checksum_part: str

# ============================================================
# CCT COMPRESSION ENGINE
# ============================================================
class CCT_Compression_Engine:
    """
    Conditional Collapse Theory Compression Engine
    Uses SuperBoolean selection + ODE trajectory + Memory Pruning
    """
    def __init__(self,
                 num_algorithms: int = 100,
                 chunk_size: int = 4096,
                 prune_threshold: float = 0.1,
                 collapse_threshold: float = 0.01):
        
        self.CHUNK_SIZE = chunk_size
        self.PRUNE_THRESHOLD = prune_threshold
        self.COLLAPSE_THRESHOLD = collapse_threshold
        
        # Initialize 100 Algorithms (Simulated subset for demo)
        self.algorithms = self._initialize_algorithm_manifold(num_algorithms)
        self.active_algos = set(range(num_algorithms))
        
        # 16-Element State Vector
        self.state_vector = np.zeros(16)
        self.entropy_trajectory = []
        self.algorithm_sequence = []
        self.total_original_size = 0
        self.total_compressed_size = 0
        
    def _initialize_algorithm_manifold(self, n: int) -> List[CompressionAlgorithm]:
        """Create the SuperBoolean Manifold of Compressors"""
        algos = []
        categories = ['dictionary', 'statistical', 'transform', 'semantic', 'hybrid']
        for i in range(n):
            algos.append(CompressionAlgorithm(
                name=f"Algo_{i}_{categories[i % len(categories)]}",
                id=i,
                category=categories[i % len(categories)],
                cost_weight=np.random.uniform(0.1, 1.0)
            ))
        return algos

    def calculate_entropy(self, data: bytes) -> float:
        """Calculate Shannon Entropy of chunk (E01)"""
        if not data: return 0.0
        freq = np.bincount(np.frombuffer(data, dtype=np.uint8))
        p = freq / np.sum(freq)
        p = p[p > 0]
        return -np.sum(p * np.log2(p))

    def superboolean_select(self, chunk: bytes) -> int:
        """
        SuperBoolean Selection: Collapse algorithm superposition based on chunk content
        """
        # 1. Analyze Chunk Semantic Signature (Simulated)
        chunk_entropy = self.calculate_entropy(chunk)
        self.state_vector[0] = chunk_entropy / 8.0  # Normalize
        
        # 2. Update Algorithm Weights (ODE Dynamics)
        # Algorithms that match the entropy profile gain weight
        for idx in self.active_algos:
            algo = self.algorithms[idx]
            # Simulate compatibility score
            compatibility = np.random.uniform(0.1, 1.0) 
            if algo.category == 'statistical' and chunk_entropy < 4.0:
                compatibility *= 1.5 # Boost for low entropy
            
            # ODE Weight Update: dw = alpha * compatibility - decay
            algo.weight += 0.01 * compatibility - 0.001 * algo.weight
            algo.weight = np.clip(algo.weight, 0, 1)
        
        # 3. Collapse to Best Algorithm
        best_algo = max(self.active_algos, key=lambda i: self.algorithms[i].weight)
        return best_algo

    def compress_chunk(self, chunk: bytes, algo_id: int) -> Tuple[bytes, float]:
        """Execute the collapsed algorithm"""
        algo = self.algorithms[algo_id]
        
        # Simulate compression using standard libs mapped to algo categories
        if 'dictionary' in algo.name:
            compressed = zlib.compress(chunk)
        elif 'statistical' in algo.name:
            compressed = bz2.compress(chunk)
        elif 'transform' in algo.name:
            compressed = lzma.compress(chunk)
        else:
            compressed = chunk # No compression
        
        ratio = len(compressed) / len(chunk) if len(chunk) > 0 else 1.0
        return compressed, ratio

    def prune_algorithms(self):
        """Entropy-Gated Forgetting of Poor Algorithms (Memory Pruning)"""
        to_prune = []
        for idx in self.active_algos:
            algo = self.algorithms[idx]
            # Prune if weight is low (low collapse potential)
            if algo.weight < self.PRUNE_THRESHOLD:
                to_prune.append(idx)
        
        for idx in to_prune:
            self.active_algos.remove(idx)
            self.algorithms[idx].active = False
            
        if to_prune:
            print(f"  [PRUNE] Removed {len(to_prune)} low-performance algorithms")

    def decompress_chunk(self, compressed_chunk: bytes, algo_id: int) -> bytes:
        """Execute the collapsed algorithm for decompression"""
        algo = self.algorithms[algo_id]

        if 'dictionary' in algo.name:
            decompressed = zlib.decompress(compressed_chunk)
        elif 'statistical' in algo.name:
            decompressed = bz2.decompress(compressed_chunk)
        elif 'transform' in algo.name:
            decompressed = lzma.decompress(compressed_chunk)
        else:
            decompressed = compressed_chunk

        return decompressed

    def compress_directory(self, dir_path: str, output_path: str, verbose: bool = True) -> Dict:
        """Compress an entire directory into a single CCT archive"""
        if verbose:
            print("=" * 70)
            print("CCT COMPRESSION ENGINE: DIRECTORY ARCHIVE")
            print("=" * 70)

        dir_path = Path(dir_path)
        if not dir_path.is_dir():
            raise ValueError(f"'{dir_path}' is not a directory")

        # Create tar archive in memory
        tar_buffer = io.BytesIO()
        with tarfile.open(fileobj=tar_buffer, mode='w') as tar:
            tar.add(dir_path, arcname=dir_path.name)
        tar_data = tar_buffer.getvalue()

        if verbose:
            print(f"  Directory archived: {len(tar_data)} bytes")
            print(f"  Compressing archive data...")

        # Compress the tar data
        self.total_original_size = len(tar_data)
        chunks = [tar_data[i:i + self.CHUNK_SIZE] for i in range(0, len(tar_data), self.CHUNK_SIZE)]
        compressed_data = b''
        chunk_metadata = []

        initial_entropy = self.calculate_entropy(tar_data)

        for t, chunk in enumerate(chunks):
            selected_algo = self.superboolean_select(chunk)
            compressed_chunk, ratio = self.compress_chunk(chunk, selected_algo)
            compressed_data += compressed_chunk
            self.total_compressed_size += len(compressed_chunk)

            chunk_metadata.append(selected_algo)
            self.algorithms[selected_algo].success_history.append(ratio)
            self.algorithms[selected_algo].weight += 0.05 * (1.0 - ratio)

            current_entropy = self.calculate_entropy(compressed_chunk)
            self.entropy_trajectory.append(current_entropy)

            if t % 10 == 0:
                self.prune_algorithms()

            if verbose and t % 50 == 0:
                print(f"  Chunk {t}: Algo={self.algorithms[selected_algo].name} | Ratio={ratio:.2f} | Active Algos={len(self.active_algos)}")

        full_signature = compressed_data + bytes(chunk_metadata)
        final_checksum = hashlib.sha256(full_signature).hexdigest()

        global_ratio = self.total_compressed_size / self.total_original_size
        entropy_reduction = (initial_entropy - np.mean(self.entropy_trajectory)) / initial_entropy

        # Save compressed data with metadata
        # Format: CHUNK_SIZE|ORIGINAL_SIZE|IS_DIR|END| + compressed_data
        metadata = f"{self.CHUNK_SIZE}|{self.total_original_size}|1|".encode()
        metadata += bytes(chunk_metadata)
        metadata += b"|END|"

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

        result = {
            "status": "COLLAPSED",
            "original_size": self.total_original_size,
            "compressed_size": self.total_compressed_size + len(metadata),
            "compression_ratio": global_ratio,
            "entropy_reduction": entropy_reduction,
            "checksum": final_checksum,
            "algorithm_sequence_length": len(chunk_metadata),
            "active_algorithms_remaining": len(self.active_algos),
            "entropy_trajectory": self.entropy_trajectory
        }

        if verbose:
            print("-" * 70)
            print(f"FINAL COMPRESSION METRICS:")
            print(f"  Original Size: {self.total_original_size} bytes")
            print(f"  Compressed Size: {self.total_compressed_size + len(metadata)} bytes")
            print(f"  Compression Ratio: {global_ratio:.4f}")
            print(f"  Entropy Reduction: {entropy_reduction * 100:.1f}%")
            print(f"  Final Checksum: {final_checksum[:16]}...")
            print(f"  Algorithms Pruned: {100 - len(self.active_algos)}")
            print("=" * 70)

        return result

    def decompress_directory(self, file_path: str, output_path: str, verbose: bool = True) -> Dict:
        """Decompress a CCT archive into a directory"""
        if verbose:
            print("=" * 70)
            print("CCT DECOMPRESSION ENGINE: DIRECTORY ARCHIVE")
            print("=" * 70)

        raw_data = Path(file_path).read_bytes()

        # Parse metadata
        metadata_end = raw_data.find(b"|END|")
        if metadata_end == -1:
            raise ValueError("Invalid CCT file: missing metadata marker")

        metadata = raw_data[:metadata_end]
        compressed_data = raw_data[metadata_end + 5:]

        parts = metadata.split(b"|", 3)
        chunk_size = int(parts[0])
        original_size = int(parts[1])
        is_directory = parts[2] == b"1"
        algo_sequence = list(parts[3])

        if not is_directory:
            raise ValueError("This is a file archive, not a directory archive")

        # Decompress chunks
        decompressed_data = b''
        offset = 0
        chunks_processed = 0

        for t, algo_id in enumerate(algo_sequence):
            remaining = original_size - len(decompressed_data)
            current_chunk_size = min(chunk_size, remaining)

            if offset >= len(compressed_data):
                break

            # Try to decompress
            decompressed = None
            for test_algo_id in self.active_algos:
                try:
                    test_chunk = compressed_data[offset:offset + current_chunk_size + 100]
                    test_decompressed = self.decompress_chunk(test_chunk, test_algo_id)
                    if len(test_decompressed) == current_chunk_size:
                        decompressed = test_decompressed
                        break
                except Exception:
                    continue

            if decompressed is None:
                try:
                    chunk = compressed_data[offset:offset + current_chunk_size]
                    decompressed = self.decompress_chunk(chunk, algo_id)
                except Exception as e:
                    if verbose:
                        print(f"  Warning: Failed to decompress chunk {t}: {e}")
                    decompressed = compressed_data[offset:offset + current_chunk_size]

            decompressed_data += decompressed
            chunks_processed += 1

            if verbose and t % 50 == 0:
                print(f"  Chunk {t}: Decompressed | Total: {len(decompressed_data)}/{original_size} bytes")

        # Verify
        final_checksum = hashlib.sha256(decompressed_data).hexdigest()
        success = len(decompressed_data) == original_size

        # Extract tar archive
        output_dir = Path(output_path)
        output_dir.mkdir(parents=True, exist_ok=True)

        tar_buffer = io.BytesIO(decompressed_data)
        with tarfile.open(fileobj=tar_buffer, mode='r') as tar:
            tar.extractall(path=output_dir)

        result = {
            "status": "RECONSTRUCTED" if success else "PARTIAL",
            "original_size": original_size,
            "decompressed_size": len(decompressed_data),
            "checksum": final_checksum,
            "chunks_processed": chunks_processed,
            "success": success
        }

        if verbose:
            print("-" * 70)
            print(f"FINAL DECOMPRESSION METRICS:")
            print(f"  Expected Size: {original_size} bytes")
            print(f"  Decompressed Size: {len(decompressed_data)} bytes")
            print(f"  Chunks Processed: {chunks_processed}")
            print(f"  Success: {success}")
            print(f"  Checksum: {final_checksum[:16]}...")
            print("=" * 70)

        return result

    def compress_file(self, file_path: str, output_path: str, verbose: bool = True) -> Dict:
        """Compress a file using CCT engine, saving metadata for decompression"""
        if verbose:
            print("=" * 70)
            print("CCT COMPRESSION ENGINE: SEMANTIC ENTROPY COLLAPSE")
            print("=" * 70)

        data = Path(file_path).read_bytes()
        self.total_original_size = len(data)

        chunks = [data[i:i + self.CHUNK_SIZE] for i in range(0, len(data), self.CHUNK_SIZE)]
        compressed_data = b''
        chunk_metadata = []

        initial_entropy = self.calculate_entropy(data)

        for t, chunk in enumerate(chunks):
            selected_algo = self.superboolean_select(chunk)
            compressed_chunk, ratio = self.compress_chunk(chunk, selected_algo)
            compressed_data += compressed_chunk
            self.total_compressed_size += len(compressed_chunk)

            chunk_metadata.append(selected_algo)
            self.algorithms[selected_algo].success_history.append(ratio)
            self.algorithms[selected_algo].weight += 0.05 * (1.0 - ratio)

            current_entropy = self.calculate_entropy(compressed_chunk)
            self.entropy_trajectory.append(current_entropy)

            if t % 10 == 0:
                self.prune_algorithms()

            if verbose and t % 50 == 0:
                print(f"  Chunk {t}: Algo={self.algorithms[selected_algo].name} | Ratio={ratio:.2f} | Active Algos={len(self.active_algos)}")

        full_signature = compressed_data + bytes(chunk_metadata)
        final_checksum = hashlib.sha256(full_signature).hexdigest()

        global_ratio = self.total_compressed_size / self.total_original_size
        entropy_reduction = (initial_entropy - np.mean(self.entropy_trajectory)) / initial_entropy

        # Save compressed data + metadata (chunk_size, original_size, algorithm sequence)
        metadata = f"{self.CHUNK_SIZE}|{self.total_original_size}|0|".encode()
        metadata += bytes(chunk_metadata)
        metadata += b"|END|"

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

        result = {
            "status": "COLLAPSED",
            "original_size": self.total_original_size,
            "compressed_size": self.total_compressed_size + len(metadata),
            "compression_ratio": global_ratio,
            "entropy_reduction": entropy_reduction,
            "checksum": final_checksum,
            "algorithm_sequence_length": len(chunk_metadata),
            "active_algorithms_remaining": len(self.active_algos),
            "entropy_trajectory": self.entropy_trajectory
        }

        if verbose:
            print("-" * 70)
            print(f"FINAL COMPRESSION METRICS:")
            print(f"  Original Size: {self.total_original_size} bytes")
            print(f"  Compressed Size: {self.total_compressed_size + len(metadata)} bytes")
            print(f"  Compression Ratio: {global_ratio:.4f}")
            print(f"  Entropy Reduction: {entropy_reduction * 100:.1f}%")
            print(f"  Final Checksum: {final_checksum[:16]}...")
            print(f"  Algorithms Pruned: {100 - len(self.active_algos)}")
            print("=" * 70)

        return result

    def decompress_file(self, file_path: str, output_path: str, verbose: bool = True) -> Dict:
        """Decompress a CCT-compressed file"""
        if verbose:
            print("=" * 70)
            print("CCT DECOMPRESSION ENGINE: RECONSTRUCTING DATA")
            print("=" * 70)

        raw_data = Path(file_path).read_bytes()

        # Parse metadata
        metadata_end = raw_data.find(b"|END|")
        if metadata_end == -1:
            raise ValueError("Invalid CCT file: missing metadata marker")

        metadata = raw_data[:metadata_end]
        compressed_data = raw_data[metadata_end + 5:]

        parts = metadata.split(b"|", 3)
        chunk_size = int(parts[0])
        original_size = int(parts[1])
        is_directory = parts[2] == b"1"
        algo_sequence = list(parts[3])

        if is_directory:
            raise ValueError("This is a directory archive, not a file archive")

        # Decompress chunks
        decompressed_data = b''
        offset = 0
        chunks_processed = 0

        for t, algo_id in enumerate(algo_sequence):
            # Calculate chunk size (last chunk may be smaller)
            remaining = original_size - len(decompressed_data)
            current_chunk_size = min(chunk_size, remaining)

            # Extract compressed chunk (we need to track sizes)
            # For simplicity, assume we stored chunk boundaries
            # In a real implementation, you'd store chunk sizes in metadata
            if offset >= len(compressed_data):
                break

            # Try to decompress - we need actual chunk boundaries
            # For now, use a simple approach: try each algorithm
            decompressed = None
            for test_algo_id in self.active_algos:
                try:
                    test_chunk = compressed_data[offset:offset + current_chunk_size + 100]
                    test_decompressed = self.decompress_chunk(test_chunk, test_algo_id)
                    if len(test_decompressed) == current_chunk_size:
                        decompressed = test_decompressed
                        break
                except Exception:
                    continue

            if decompressed is None:
                # Fallback: use the algorithm ID from metadata
                try:
                    chunk = compressed_data[offset:offset + current_chunk_size]
                    decompressed = self.decompress_chunk(chunk, algo_id)
                except Exception as e:
                    if verbose:
                        print(f"  Warning: Failed to decompress chunk {t}: {e}")
                    decompressed = compressed_data[offset:offset + current_chunk_size]

            decompressed_data += decompressed
            chunks_processed += 1

            if verbose and t % 50 == 0:
                print(f"  Chunk {t}: Decompressed | Total: {len(decompressed_data)}/{original_size} bytes")

        # Verify
        final_checksum = hashlib.sha256(decompressed_data).hexdigest()
        success = len(decompressed_data) == original_size

        # Write output
        with open(output_path, 'wb') as f:
            f.write(decompressed_data)

        result = {
            "status": "RECONSTRUCTED" if success else "PARTIAL",
            "original_size": original_size,
            "decompressed_size": len(decompressed_data),
            "checksum": final_checksum,
            "chunks_processed": chunks_processed,
            "success": success
        }

        if verbose:
            print("-" * 70)
            print(f"FINAL DECOMPRESSION METRICS:")
            print(f"  Expected Size: {original_size} bytes")
            print(f"  Decompressed Size: {len(decompressed_data)} bytes")
            print(f"  Chunks Processed: {chunks_processed}")
            print(f"  Success: {success}")
            print(f"  Checksum: {final_checksum[:16]}...")
            print("=" * 70)

        return result


def main():
    parser = argparse.ArgumentParser(
        description="CCT Compression Engine - Compress/Decompress binary files and directories",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python ai_compression.py compress input.bin output.cct
  python ai_compression.py compress ./mydir output.cct
  python ai_compression.py decompress input.cct output.bin
  python ai_compression.py decompress archive.cct ./restored_dir
  python ai_compression.py compress input.bin output.cct --chunk-size 8192 --quiet
        """
    )

    parser.add_argument('mode', choices=['compress', 'decompress'],
                        help='Operation mode: compress or decompress')
    parser.add_argument('input', help='Input file or directory path')
    parser.add_argument('output', help='Output file path')
    parser.add_argument('--chunk-size', type=int, default=4096,
                        help='Chunk size in bytes (default: 4096)')
    parser.add_argument('--algorithms', type=int, default=100,
                        help='Number of algorithms to initialize (default: 100)')
    parser.add_argument('--prune-threshold', type=float, default=0.1,
                        help='Pruning threshold (default: 0.1)')
    parser.add_argument('--quiet', action='store_true',
                        help='Suppress verbose output')

    args = parser.parse_args()

    # Validate input exists
    input_path = Path(args.input)
    if not input_path.exists():
        print(f"Error: Input path '{args.input}' not found")
        return 1

    # Initialize engine
    engine = CCT_Compression_Engine(
        num_algorithms=args.algorithms,
        chunk_size=args.chunk_size,
        prune_threshold=args.prune_threshold
    )

    try:
        if args.mode == 'compress':
            if input_path.is_dir():
                engine.compress_directory(args.input, args.output, verbose=not args.quiet)
            else:
                engine.compress_file(args.input, args.output, verbose=not args.quiet)
        else:
            # For decompression, we need to detect if it's a directory or file archive
            raw_data = input_path.read_bytes()
            metadata_end = raw_data.find(b"|END|")
            if metadata_end == -1:
                print("Error: Invalid CCT file format")
                return 1

            metadata = raw_data[:metadata_end]
            parts = metadata.split(b"|", 3)
            is_directory = parts[2] == b"1"

            if is_directory:
                engine.decompress_directory(args.input, args.output, verbose=not args.quiet)
            else:
                engine.decompress_file(args.input, args.output, verbose=not args.quiet)
        return 0
    except Exception as e:
        print(f"Error: {e}")
        return 1


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