from os import path
from pathlib import Path

import os
import string
import subprocess
import glob

MSVC_VS_PATTERN = "Program Files*/Microsoft Visual Studio/20*/*/VC/Tools/MSVC/*/"

# 64 bit
MSVC_COMPILER_PATTERN = "bin/Hostx64/x64/cl.exe"
MSVC_INCLUDE_PATTERN = "include/"
MSVC_LIB_PATTERN = "lib/x64/"


def find_msvc():
    potential_msvcs = []
    available_drives = ['%s:' % d for d in string.ascii_uppercase if os.path.exists('%s:' % d)]
    for available_drive in available_drives:
        glob_pattern = available_drive + "/" + MSVC_VS_PATTERN
        potential_msvcs += glob.glob(glob_pattern)

    if not potential_msvcs:
        raise RuntimeError("Could not locate MSVC")

    msvc = Path(potential_msvcs[0])

    msvc_compiler = msvc / MSVC_COMPILER_PATTERN
    msvc_include = msvc / MSVC_INCLUDE_PATTERN
    msvc_lib = msvc / MSVC_LIB_PATTERN

    if not path.exists(msvc_compiler):
        raise RuntimeError("Could not locate MSVC compiler")

    if not path.exists(msvc_include):
        raise RuntimeError("Could not locate MSVC includes")

    if not path.exists(msvc_lib):
        raise RuntimeError("Could not locate MSVC libs")

    return msvc_compiler, msvc_include, msvc_lib


WINDOWS_DEVKIT_LOCATION = 'Program Files*/Windows Kits/10/Include/*/'


def find_windows_dev_kit():
    potential_dev_kits = []

    available_drives = ['%s:' % d for d in string.ascii_uppercase if os.path.exists('%s:' % d)]
    for available_drive in available_drives:
        glob_pattern = available_drive + "/" + WINDOWS_DEVKIT_LOCATION
        potential_dev_kits += glob.glob(glob_pattern)

    if not potential_dev_kits:
        raise RuntimeError("Could not locate a Windows Development Kit")

    win_dev_kit_include_root = Path(potential_dev_kits[0])
    win_dev_kit_include = win_dev_kit_include_root / 'ucrt'
    win_dev_version = win_dev_kit_include_root.name
    win_dev_kit_lib_root = win_dev_kit_include_root.parents[1] / 'Lib' / win_dev_version
    win_dev_kit_lib_um = win_dev_kit_lib_root / 'um' / 'x64'
    win_dev_kit_lib_ucrt = win_dev_kit_lib_root / 'ucrt' / 'x64'

    return [win_dev_kit_include], [win_dev_kit_lib_um, win_dev_kit_lib_ucrt]


def compile_msvc(filenames: list, c_compiler: str = None, c_includes: list = None, c_libs: list = None,
                 c_compile_options: list = None, c_linker_options: list = None, mitigations: bool = False):
    c_includes = c_includes if c_includes else []
    c_libs = c_libs if c_libs else []
    c_compile_options = c_compile_options if c_compile_options else []
    c_linker_options = c_linker_options if c_linker_options else []

    # add msvc
    if not c_compiler:
        (c_compiler, msvc_includes, msvc_libs) = find_msvc()
        c_includes.append(msvc_includes)
        c_libs.append(msvc_libs)

    # add windows dev kit
    (windows_include, windows_lib) = find_windows_dev_kit()
    c_includes += windows_include
    c_libs += windows_lib

    if type(filenames) is not list:
        filenames = [filenames]

    print("Compiling source files...")

    for filename in filenames:
        print("Compiling: " + str(filename))

        command_list = [str(c_compiler),
                        "/LD",
                        "/MD",
                        "/Z7",
                        "/FAcs",
                        "/nologo",
                        "/W4",
                        "/O2"
                        ]
        if mitigations:
            command_list += ["/Qspectre"]

        for c_include in c_includes:
            command_list += ['/I', str(c_include)]

        command_list += c_compile_options
        command_list += [filename, "ex_main.c", '/link']

        for c_lib in c_libs:
            command_list += ['/LIBPATH:' + str(c_lib)]

        command_list += c_linker_options

        result = subprocess.run(command_list)
        if result.stderr:
            print("Finished compiling with errors : " + str(result.stderr))
        else:
            print("Finished compiling: " + str(filename))

    print("Finished compiling source files")


def compile_paul():
    # compile all originals with mitigations
    home = "TODO"
    pauls_home = home + "BCB\\Paul Kocher\\"
    cheangs_home = home + "BCB\\Cheang et al\\"

    files = [
        cheangs_home + "Example CV\\exCV.c",
    ]

    # compile_msvc(files)
    compile_msvc(files, mitigations=False)


if __name__ == "__main__":
    compile_paul()