import os
import tarfile
import gzip
import shutil
import argparse
import logging

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")


def extract_tar(file_path, output_dir):
    """Extract a tar archive with auto-detection of compression."""
    try:
        # "r:*" auto-detects the compression type.
        with tarfile.open(file_path, "r:*") as tar:
            tar.extractall(path=output_dir)
        logging.info(f"Extracted {file_path} to {output_dir}")
    except Exception as e:
        logging.error(f"Error extracting {file_path}: {e}")


def extract_nested_archives(root_archives, working_dir):
    """Extract multiple root .tar.gz files and all nested .tar.gz files."""
    if not os.path.exists(working_dir):
        os.makedirs(working_dir)

    for archive in root_archives:
        extract_tar(archive, working_dir)

    # Extract all nested .tar.gz files
    for root, _, files in os.walk(working_dir):
        for file in files:
            if file.endswith(".tar.gz"):
                nested_tar_path = os.path.join(root, file)
                extract_tar(nested_tar_path, root)


def extract_db2_gz_files(working_dir, dest_folder):
    """Extract all .db2.gz files and move extracted .db2 files to the destination folder."""
    if not os.path.exists(dest_folder):
        os.makedirs(dest_folder)

    for root, _, files in os.walk(working_dir):
        for file in files:
            if file.endswith(".db2.gz"):
                db2_gz_path = os.path.join(root, file)
                db2_path = db2_gz_path[:-3]  # Remove .gz extension

                # Extract .db2.gz
                try:
                    with gzip.open(db2_gz_path, "rb") as f_in:
                        with open(db2_path, "wb") as f_out:
                            shutil.copyfileobj(f_in, f_out)
                    os.remove(db2_gz_path)  # Remove the original .db2.gz file after extraction
                    logging.info(f"Extracted: {db2_gz_path} -> {db2_path}")
                except Exception as e:
                    logging.error(f"Error extracting {db2_gz_path}: {e}")

                # Move extracted .db2 file
                shutil.move(db2_path, os.path.join(dest_folder, os.path.basename(db2_path)))

    logging.info(f"All .db2 files moved to {dest_folder}")


def main():
    parser = argparse.ArgumentParser(description="Extract nested tar.gz downloaded from TLDR.")
    parser.add_argument("--root-archives", nargs="+", required=True, help="Paths to the root .tar.gz files (space-separated for multiple files)")
    parser.add_argument("--working-dir", default="extract_working_dir", help="Temporary extraction folder")
    parser.add_argument("--dest-folder", default="db2_files", help="Destination folder for .db2 files")
    args = parser.parse_args()

    logging.info("Starting extraction process...")

    extract_nested_archives(args.root_archives, args.working_dir)
    extract_db2_gz_files(args.working_dir, args.dest_folder)

    logging.info("Extraction and organization complete.")


if __name__ == "__main__":
    main()
