import os
import shutil
import tempfile
import tarfile
import zipfile
import gzip
import bz2
import io
import magic

# Max bytes to inspect
MAGIC_BYTES_N = 2048

# Map MIME type to handler
HANDLERS = {
    'application/zip': 'zip',
    'application/x-tar': 'tar',
    'application/gzip': 'gz',
    'application/x-bzip2': 'bz2'
}

def detect_mime(path):
    """Detect the MIME type using file magic."""
    mime = magic.Magic(mime=True)
    return mime.from_file(path)

def decompress_tar(path, tempdir):
    with tarfile.open(path, mode='r:*') as tar:
        tar.extractall(tempdir)
        return [os.path.join(tempdir, member.name) for member in tar.getmembers() if member.isfile()]

def decompress_zip(path, tempdir):
    with zipfile.ZipFile(path, 'r') as zf:
        zf.extractall(tempdir)
        return [os.path.join(tempdir, name) for name in zf.namelist()]

def decompress_gz(path, tempdir):
    out_path = os.path.join(tempdir, os.path.basename(path).replace('.gz', ''))
    with gzip.open(path, 'rb') as gz_in:
        with open(out_path, 'wb') as out_f:
            shutil.copyfileobj(gz_in, out_f)
    return [out_path]

def decompress_bz2(path, tempdir):
    out_path = os.path.join(tempdir, os.path.basename(path).replace('.bz2', ''))
    with bz2.open(path, 'rb') as bz2_in:
        with open(out_path, 'wb') as out_f:
            shutil.copyfileobj(bz2_in, out_f)
    return [out_path]

def decompress_by_type(path, mime_type, tempdir):
    if mime_type == 'application/zip':
        return decompress_zip(path, tempdir)
    elif mime_type == 'application/x-tar':
        return decompress_tar(path, tempdir)
    elif mime_type == 'application/gzip':
        return decompress_gz(path, tempdir)
    elif mime_type == 'application/x-bzip2':
        return decompress_bz2(path, tempdir)
    else:
        return []

def recursive_decompress(input_path):
    """
    Recursively decompress a file until no known compression is detected.
    Returns list of final uncompressed file paths.
    """
    current_files = [input_path]
    final_outputs = []

    while current_files:
        next_files = []
        for file in current_files:
            mime = detect_mime(file)
            if mime in HANDLERS:
                print(f"[+] Decompressing {file} ({mime})")
                with tempfile.TemporaryDirectory() as tempdir:
                    decompressed = decompress_by_type(file, mime, tempdir)
                    next_files.extend(decompressed)
            else:
                final_outputs.append(file)

        current_files = next_files

    return final_outputs


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Recursively decompress compressed files using magic bytes.")
    parser.add_argument("file", help="Path to the compressed file")
    args = parser.parse_args()

    if not os.path.isfile(args.file):
        print("File does not exist.")
        exit(1)

    final_files = recursive_decompress(args.file)
    print("\n[✓] Final decompressed file(s):")
    for f in final_files:
        print(f"  - {f}")