#!/usr/bin/env python3

import fnmatch
import subprocess
import os
import random
import string
import re
import click
from progressbar import ProgressBar

FILE_EXTENSIONS = ["cpp", "c", "h", "hpp", "cc", "hh", "cxx", "hxx"]
INCLUDE_REGEXP = r'(#include\s*[<"])([^>]*)([>"])'
RANDOM_STRING = ''.join(random.choice(string.ascii_uppercase) for _ in range(78))
PLACEHOLDER_COMMENT = str.encode("//{}\n".format(RANDOM_STRING))

def run_build(build_command):
    """Returns False if the build command failed"""
    try:
        with open(os.devnull, 'w') as devnull:
            subprocess.check_call(build_command, shell=True, stdout=devnull, stderr=devnull)
    except subprocess.CalledProcessError:
        return False
    return True

def print_status(status, color, include_directive):
    print(" {} {}:{}".format(
        click.style("{:>7}".format(status), color),
        os.path.relpath(include_directive[0]), include_directive[1] + 1
    ))

@click.command()
@click.version_option(version="1.0.0")
@click.option('--build_command', help='The build command to check if an include is needed.')
def main(build_command):
    if not build_command:
        build_command = click.prompt("Enter build command")
    print("Testing build command ... ", end="", flush=True)
    if not run_build(build_command):
        click.secho("Failed", fg='red')
        exit(1)

    click.secho("OK", fg='green')
    source_files = []

    print("Scanning for C/C++ source files ... ", end="", flush=True)
    for extension in FILE_EXTENSIONS:
        for root, _, filenames in os.walk('.'):
            for filename in fnmatch.filter(filenames, '*.' + extension):
                path = os.path.join(root, filename)
                if 'wx' in path:
                    print(path)
                source_files.append(path)
    click.secho("OK", fg='green')

    include_directives = []

    print("Scanning for include directives ... ", end="", flush=True)
    for filename in source_files:
        with open(filename, "r", errors='ignore') as file:
            for linenumber, line in enumerate(file.readlines()):
                if line == RANDOM_STRING:
                    click.secho("\nFatal error: {} contains {}.".format(
                        click.format_filename(filename), RANDOM_STRING
                    ), fg='red')
                match = re.search(INCLUDE_REGEXP, line)
                if match:
                    include_directives.append((filename, linenumber,))
    click.secho("OK", fg='green')

    random.shuffle(include_directives)
    click.secho("Testing which include directives can be removed", bold=True)
    progress_bar = ProgressBar(redirect_stdout=True)
    n_removed = 0
    n_skipped = 0

    for include_directive in progress_bar(include_directives):
        linenumber = include_directive[1]
        lines = []
        with open(include_directive[0], "rb") as file:
            lines = list(file.readlines())
        line_backup = lines[linenumber]
        lines[linenumber] = str.encode(
            re.sub(INCLUDE_REGEXP, r"\1{}\3".format(RANDOM_STRING), line_backup.decode('utf-8'))
        )
        try:
            with open(include_directive[0], "wb") as file:
                file.writelines(lines)
        except PermissionError:
            print_status("Error", "red", include_directive)
            n_skipped += 1
            continue
        if run_build(build_command):
            # Include directive doesn't influence the build, undo change and skip it
            lines[linenumber] = line_backup
            with open(include_directive[0], "wb") as file:
                file.writelines(lines)
            print_status("Skipped", "yellow", include_directive)
            n_skipped += 1
            continue
        lines[linenumber] = str.encode("//{}\n".format(RANDOM_STRING))
        with open(include_directive[0], "wb") as file:
            file.writelines(lines)
        if run_build(build_command):
            print_status("Removed", "green", include_directive)
            n_removed += 1
        else:
            lines[linenumber] = line_backup
            with open(include_directive[0], "wb") as file:
                file.writelines(lines)
            print_status("Needed", "blue", include_directive)

    for filename in source_files:
        with open(filename, "rb") as file:
            lines = [line for line in file.readlines() if line != PLACEHOLDER_COMMENT]
        with open(filename, "wb") as file:
            file.writelines(lines)

    click.secho("Removed {} of {} includes ({} skipped).".format(
        click.style(str(n_removed), fg='green', bold=True),
        click.style(str(len(include_directives) - n_skipped), bold=True),
        click.style(str(n_skipped), fg='yellow', bold=True),
    ))

if __name__ == '__main__':
    main()
