#!/usr/bin/env python3

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

FILE_EXTENSIONS = ["cpp", "c", "h", "hpp", "cc", "hh", "cxx", "hxx"]
INCLUDE_REGEXP = r'(#\s*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.2.0")
@click.option('--build_command', help='The build command to check if an include is needed.')
@click.option('--exclude_files',
              help='Exclude any source file which name contains this regular expression.')
@click.option('--exclude_includes',
              help='Exclude any include directive which contains this regular expression.',
              default=r'// IWYU pragma: keep', show_default=True)
def main(build_command, exclude_files, exclude_includes):
    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 = []

    exclude_count = 0
    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(os.path.relpath(root, '.'), filename)
                if exclude_files and re.search(exclude_files, path):
                    exclude_count += 1
                else:
                    source_files.append(path)
    click.secho("OK ({} excluded)".format(exclude_count), 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:
                    if not re.search(exclude_includes, line):
                        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
    should_quit = False

    def handle_sigint(_signal, _frame):
        nonlocal should_quit
        should_quit = True
    signal.signal(signal.SIGINT, handle_sigint) # Remove placeholders on Ctrl+C

    for include_directive in progress_bar(include_directives):
        if should_quit:
            break
        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{}\2".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) or should_quit:
            # Include directive doesn't influence the build, undo change and skip it. If should_quit
            # is set, the user pressed Ctrl+C and the build was canceled. run_build will return
            # False, although it might have succeeded. Also undo the change in that case.
            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:
            old_lines = file.readlines()
            lines = [line for line in old_lines if line != PLACEHOLDER_COMMENT]
        if len(lines) < len(old_lines):
            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(max_content_width=100)
