import os
import gzip
import re
import traceback
from concurrent.futures import ProcessPoolExecutor, as_completed, wait
import subprocess
import time

def check_disk_space(path, min_space_gb=20, check_interval=300):
    while True:
        try:
            result = subprocess.run(['df', '-BG', path], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            if result.returncode == 0:
                output = result.stdout.split('\n')[1].split()
                available_space_gb = int(output[3].replace('G', ''))
                if available_space_gb >= min_space_gb:
                    break
                else:
                    print(f"Waiting for more than {min_space_gb}GB free space on {path}. Current available space: {available_space_gb}GB")
                    time.sleep(check_interval)
            else:
                print("Error checking disk space.")
                time.sleep(check_interval)
        except Exception as e:
            print(f"An error occurred while checking disk space: {e}")
            time.sleep(check_interval)

def extract_urls_from_file(file_path):
    urls = []
    try:
        with gzip.open(file_path, 'rt', encoding='latin-1') as file:
            for line in file:
                url_pattern = re.compile(r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+')
                line_urls = re.findall(url_pattern, line)
                urls.extend(line_urls)
    except (gzip.BadGzipFile, EOFError) as e:
        print(f"Error while reading the compressed file '{file_path}': {e}")
    except Exception as e:
        print(f"An unexpected error occurred while processing '{file_path}': {e}")
        print("Full traceback:")
        traceback.print_exc()

    return urls

def process_file(file_path):
    try:
        print(f"Processing file: {file_path}")

        # Extract URLs from the gzipped file
        urls = extract_urls_from_file(file_path)

        # Create the output file path with '_urls.txt' extension
        output_file_path = os.path.splitext(file_path)[0] + '_urls.txt'
        print(f"Output file path: {output_file_path}")

        # Write the URLs to the output file
        with open(output_file_path, 'w') as output_file:
            output_file.write('\n'.join(urls))
        print(f"URLs written to {output_file_path}")

        # Use zstd command-line tool for compression
        compressed_file_path = f'{output_file_path}.zst'
        command_compress = f'zstd -12 --long --force {output_file_path} -o {compressed_file_path}'

        try:
            # Run the compression command synchronously and wait for it to complete
            result = subprocess.run(command_compress, shell=True, check=True)
            if result.returncode == 0:
                print(f"Compressed file saved as '{compressed_file_path}'")
            else:
                print(f"Compression failed for '{output_file_path}'")
        except subprocess.CalledProcessError as compression_error:
            print(f"Compression failed for '{output_file_path}': {compression_error}")

        # Remove the original gzipped file
        os.remove(file_path)
        print(f"Original file removed: {file_path}")

        # Remove the original _urls.txt file
        os.remove(output_file_path)
        print(f"Original file removed: {output_file_path}")

        # Remove the line containing the filename (without "_urls.txt") from urls_to_download.txt
        filename = os.path.basename(output_file_path).replace('_urls.txt', '')
        command = f'sed -i "/{filename}/d" "urls_to_download.txt"'
        result = subprocess.run(command, shell=True)
        if result.returncode == 0:
            print(f"File {filename} has been successfully removed from urls_to_download.txt")
            with open('urls_to_download.txt', 'r') as file:
                remaining_count = sum(1 for line in file)
            print(f"URLs remaining to be processed: {remaining_count}")
        else:
            print(f"Failed to remove {filename} from urls_to_download.txt")

    except Exception as e:
        print(f"Error during processing {file_path}: {e}")
        traceback.print_exc()

def download_and_process_file(url):
    try:
        command = f'axel -q -n 1 {url}'
        result = subprocess.run(command, shell=True, check=True)
        if result.returncode == 0:
            file_path = os.path.join(os.getcwd(), os.path.basename(url))
            process_file(file_path)
        else:
            print(f"Download failed for {url}")

    except Exception as e:
        print(f"Error during download and processing {url}: {e}")

def main():
    check_disk_space('/dev/sda1')

    with open('urls_to_download.txt', 'r') as file:
        urls = file.readlines()

    urls = [url.strip() for url in urls]

    download_concurrency_level = 7

    with ProcessPoolExecutor(max_workers=download_concurrency_level) as executor:
        print("Submitting tasks to the ProcessPoolExecutor...")
        futures = [executor.submit(download_and_process_file, url) for url in urls]
        print(f"Submitted {len(futures)} tasks.")

        print("Waiting for tasks to complete...")
        completed_futures, _ = wait(futures)
        print(f"{len(completed_futures)} tasks completed.")

        for completed_future in completed_futures:
            try:
                result = completed_future.result()
                print(f"Task result: {result}")
                # Process the result if needed
            except Exception as e:
                print(f"Error in processing future: {e}")

if __name__ == "__main__":
    main()