import os
import tarfile
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor

ARCHIVES_DIR = Path("/root/ankahi_data/torgo_archives")
EXTRACT_DIR = Path("/root/ankahi_data/raw/audio/torgo")

def extract_archive(filepath):
    print(f"Extracting {filepath.name} to {EXTRACT_DIR}...")
    try:
        with tarfile.open(filepath, "r:bz2") as tar:
            # Safely extract all files into the target directory
            # (which creates the required F01, M01 structure)
            tar.extractall(path=EXTRACT_DIR)
        print(f"Finished extracting {filepath.name}")
    except Exception as e:
        print(f"Error extracting {filepath.name}: {e}")

if __name__ == "__main__":
    if not ARCHIVES_DIR.exists():
        print(f"Directory {ARCHIVES_DIR} not found. Please run the download script first.")
        exit(1)
        
    EXTRACT_DIR.mkdir(parents=True, exist_ok=True)
    archives = list(ARCHIVES_DIR.glob("*.tar.bz2"))
    
    if not archives:
        print(f"No .tar.bz2 files found in {ARCHIVES_DIR}.")
        exit(1)
        
    print(f"Found {len(archives)} archives. Starting parallel extraction (~18 GB uncompressed)...")
    with ThreadPoolExecutor(max_workers=4) as executor:
        executor.map(extract_archive, archives)
    print(f"Extraction complete! Data is ready in {EXTRACT_DIR}.")