#!python
import argparse
import glob
import os
import yaml

# Colors
class colors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'
    PURPLE = '\033[95m'

# Script class
class Script:
    def __init__(self, *kargs):
        self.scripts = kargs

    def run(self, *args):
        for script in self.scripts:
            ret = os.system(script.format(*args))

            # TODO: print error
            if (ret != 0):
                exit(-1)

# Build class
class Build:
    cache = '.smake'
    bdir = '.smake/builds'
    tdir = '.smake/targets'

    # Default compiler and standard
    default_compiler = 'g++'
    default_standard = 'c++11'

    # TODO: each build should have a name
    # Build is build name
    def __init__(self, target, build, compiler, sources,
        idirs = [], libs = [], flags = []):
        # Set immediate properties
        self.target = target
        self.build = build
        self.sources = sources
        self.compiler = compiler

        # TODO: make function for directory creation
        # Create the cache directory
        if not os.path.exists(self.cache):
            os.mkdir(self.cache)

        # Create build dir
        if not os.path.exists(self.bdir):
            os.mkdir(self.bdir)
        
        # Create the output directory
        self.odir = self.bdir + '/' + self.build
        if not os.path.exists(self.odir):
            os.mkdir(self.odir)
        
        # Create target build directory
        self.tpath = self.tdir + '/' + self.target
        if not os.path.exists(self.tdir):
            os.mkdir(self.tdir)

        # Includes
        # TODO: filter includes, check if they exist
        self.idirs = ''
        if len(idirs) > 0:
            self.idirs = ''.join([' -I ' + idir for idir in idirs])

        # Libraries
        self.libs = ''
        if len(libs) > 0:
            self.libs = ''.join([' -l' + lib for lib in libs])

        # Flags
        self.flags = ''
        if len(flags) > 0:
            self.flags = ' '.join(flags)
        else:
            self.flags = '-std={}'.format(self.default_standard)
        
    # Update target name
    def set_target(self, target):
        self.target = target
        self.tpath = self.tdir + '/' + self.target

    # Compile a single file
    def compile(self, file, verbose = False):
        # Check if file exists
        if not os.path.exists(file):
            print(colors.FAIL + f'\tFile {file} does not exist' + colors.ENDC)
            return ''

        # Get file name without directory
        fname = os.path.basename(file)

        # Create the output file
        ofile = os.path.join(self.odir, fname.replace('.cpp',
            '.o').replace('.c', '.o'))

        # Check if compilation is necessary
        file_t = os.path.getmtime(file)
        if os.path.exists(ofile):
            ofile_t = os.path.getmtime(ofile)
            if ofile_t > file_t:
                print(colors.OKCYAN + 'Source already compiled' + colors.ENDC)
                return ofile

        # Compile the source
        cmd = '{} {} -c -o {} {} {}'.format(self.compiler, self.flags,
                ofile, file, self.idirs)

        # Print command if verbose
        if verbose:
            print(colors.PURPLE + cmd + colors.ENDC)
        else:
            print()

        # Run command and check if it was successful
        ret = os.system(cmd)
        if ret != 0:
            return ''

        # Return object
        return ofile

    # Compile all sources
    def run(self, verbose = False):
        # List of compiled files
        compiled = []

        # Compilation failure flag
        failed = False

        # Failed sources
        fsources = []

        # Generate format string
        slen = str(len(self.sources))
        fstr = colors.OKCYAN + '[{:>' + str(len(slen)) + '}/' + slen + '] ' \
            + colors.OKGREEN + 'Compiling {}... ' + colors.ENDC

        # Compile the sources
        for i in range(len(self.sources)):
            # Print the message
            print(fstr.format(i + 1, self.sources[i]), end='')

            # Compile the file (or attempt to)
            ofile = self.compile(self.sources[i], verbose)

            # Check failure
            if len(ofile) == 0:
                fsources.append(self.sources[i])
                failed = True

            # Add to compiled list
            compiled.append(ofile)

        # Check if any of the sources failed to compile
        if failed:
            # Print message
            print('\n' + colors.FAIL + 'Failed to compile the' + \
                ' following sources,' + ' skipping linking process' + \
                colors.ENDC)

            # Print failed sources
            for file in fsources:
                print(colors.PURPLE + '\t' + file + colors.ENDC)

            # Return empty for failure
            return ''
        else:
            # Command generation
            compiled_str = ' '.join(compiled)
            cmd = f'{self.compiler} {compiled_str} -o {self.tpath} {self.libs}'

            # Print message
            print(colors.OKBLUE + f'Linking executable {self.target}... '.ljust(50)
                + colors.ENDC, end='')

            # Print command if verbose
            if verbose:
                print(colors.PURPLE + cmd + colors.ENDC)
            else:
                print()

            # Check return of linking
            ret = os.system(cmd)
            if ret != 0:
                # Print message and return empty
                print(colors.FAIL + f'Failed to link target {self.target}' + \
                        colors.ENDC)

                return ''

        # Return the location of the target
        return self.tpath

# Target class
class Target:
    def __init__(self, name, modes, builds, postbuilds):
        self.name = name
        self.modes = modes
        self.builds = builds
        self.postbuilds = postbuilds

    def run(self, mode = 'default', verbose = False):
        # Empty is always a valid mode, but `default should be used'
        if len(mode) == 0:
            mode = 'default'
        
        # Retrieve run attributes
        build = self.builds[mode]

        # Run the build
        target = build.run(verbose)

        # Run postbuild with target argument, if present
        if mode in self.postbuilds:
            if len(target) == 0:
                print(colors.FAIL + '\nFailed to compile target, skipping postbuild script' + colors.ENDC)
                return

            postbuild = self.postbuilds[mode]
            print(colors.OKBLUE + '\nSucessfully compiled target, ' + \
                'running postbuild script\n' + colors.ENDC)
            postbuild.run(target)

# Global helper functions
def split_plain(elem):
    prop = elem
    if isinstance(elem, str):
        prop = prop.split(', ')
    return prop

def split(d, pr, defns):
    prop = split_plain(d[pr])

    out = []
    for i in range(len(prop)):
        if prop[i] in defns:
            value = defns[prop[i]]

            if isinstance(value, list):
                out.extend(value)
            else:
                out.append(value)
        else:
            out.append(prop[i])
    
    return out

def concat(ldicts):
    out = {}
    for d in ldicts:
        out.update(d)
    return out

# Config class
class Config:
    # Constructor takes no argument
    def __init__(self):
        # Initialize targets to empty
        self.targets = {}

        # Get config file from current dir
        if os.path.exists('smake.yaml'):
            self.load_file('smake.yaml')

    # Reads definitions from variables like sources, includes, etc.
    def load_definitions(self, smake):
        # TODO: error on duplicate definition

        # Output dictionary
        defns = {}
        
        # Load definitions
        if 'definitions' in smake:
            for dgroup in smake['definitions']:
                key, value = next(iter(dgroup.items()))
                value = split_plain(value)
                defns.update({key: value})
        
        # Default compiler and standard
        if 'default_compiler' in smake:
            Build.default_compiler = smake['default_compiler']
        
        if 'default_standard' in smake:
            Build.default_standard = smake['default_standard']
        
        # Return the dictionary
        return defns
    
    # Create a build object
    def load_build(self, build, defns):
        name = list(build)[0]
        properties = {}
        for d in build[name]:
            properties.update(d)

        # Preprocess properties
        sources = split(properties, 'sources', defns)

        # Optional properties
        includes = []
        libraries = []
        flags = []
        compiler = Build.default_compiler

        if 'includes' in properties:
            includes = split(properties, 'includes', defns)
        
        if 'libraries' in properties:
            libraries = split(properties, 'libraries', defns)
        
        if 'flags' in properties:
            flags = split(properties, 'flags', defns)
        
        if 'compiler' in properties:
            compiler = properties['compiler']

        # Create and return the object
        return Build('smake-build', name, compiler, sources, includes,
            libraries, flags)

    def load_all_builds(self, smake, defns):
        # Check that builds actually exists
        if 'builds' not in smake:
            print(colors.FAIL + 'No builds defined in smake.yaml' + colors.ENDC)
            exit(-1)

        blist = {}
        for b in smake['builds']:
            name = list(b)[0]
            build = self.load_build(b, defns)
            blist.update({name: build})
        
        return blist

    def load_target(self, target, blist, defns):
        name = list(target)[0]
        properties = {}
        for d in target[name]:
            properties.update(d)
        
        # Preprocess properties
        modes = split(properties, 'modes', defns)

        # Gets builds and postbuilds
        builds = concat(properties['builds'])

        postbuilds = {}
        if 'postbuild' in properties:
            postbuilds = concat(properties['postbuild'])

        # Preprocess predefined things
        # TODO: separate methods
        for b in builds:
            bname = builds[b]

            if bname in blist:
                builds[b] = blist[bname]
                builds[b].set_target(name)
            # TODO: errir handling here
        
        # If the postbuild is a string, then convert to Script
        for pe in postbuilds:
            pname = postbuilds[pe]

            if pname in defns:
                postbuilds[pe] = defns[pname]
            else:
                postbuilds[pe] = Script(pname)

        return Target(name, modes, builds, postbuilds)

    def load_all_targets(self, smake, builds, defns):
        # Check that targets exist
        if 'targets' not in smake:
            print(colors.FAIL + 'No targets specified in smake.yaml' + colors.ENDC)
            exit(-1)

        tlist = {}
        for t in smake['targets']:
            name = list(t)[0]
            target = self.load_target(t, builds, defns)
            tlist.update({name: target})

        return tlist

    # Read config file
    def load_file(self, file):
        # Open and read the config
        with open(file, 'r') as file:
            smake = yaml.safe_load(file)
        
        # Load the definitions
        defns = self.load_definitions(smake)

        # Load all builds
        builds = self.load_all_builds(smake, defns)

        # Load all targets
        self.targets = self.load_all_targets(smake, builds, defns)

    # List all targets
    def list_targets(self):
        # Return if no targets
        if len(self.targets) == 0:
            print(colors.FAIL + 'No targets found' + colors.ENDC)
            return

        # Max string length of target name
        maxlen = 0
        for t in self.targets:
            if len(t) > maxlen:
                maxlen = len(t)
        
        # Padding
        maxlen += 5

        # Header message
        fmt = colors.OKCYAN + '{:<' + str(maxlen) + '}{}' + colors.ENDC
        print(fmt.format('Target', 'Modes'))

        # Print the targets
        for t in self.targets:
            modes = ''
            for i in range(len(self.targets[t].modes)):
                modes += self.targets[t].modes[i]

                if i != len(self.targets[t].modes) - 1:
                    modes += ', '
            
            fmt = colors.OKBLUE + '{:<' + str(maxlen) + '}' + \
                colors.PURPLE + '{}' + colors.ENDC
            print(fmt.format(t, modes))
    
    # Run the correct target
    def run(self, target, mode = 'default', verbose = False):
        # If the target is all, then run all targets
        if target == 'all':
            for t in self.targets:
                self.targets[t].run(mode, verbose)
            return

        # Check if the target is valid
        if target in self.targets:
            self.targets[target].run(mode, verbose)
        else:
            print(colors.FAIL + f'No target {target} found.' + \
                ' Perhaps you meant one of the following:' + colors.ENDC)
            for valid in list(self.targets.keys()):
                print(colors.PURPLE + f'\t{valid}' + colors.ENDC)

# Build the parser
parser = argparse.ArgumentParser()

# {TARGET} -m {EXECUTOR} -j{THREADS}
parser.add_argument("target", help="Target name", nargs='?', default='all')
parser.add_argument("-m", "--mode", help="Execution mode", default='')
parser.add_argument("-j", "--threads",
                    help="Number of concurrent threads", type=int, default=8)

# Special args
parser.add_argument("-l", "--list", help="List all targets", action='store_true')
parser.add_argument("--clear-cache", help="Clear smake cache", action='store_true')
parser.add_argument("-v", "--verbose", help="Verbose mode", action='store_true')

# Read the arguments
args = parser.parse_args()

# Create the local config
config = Config()

# Run the target
if args.list:
    config.list_targets()
elif args.clear_cache:
    os.system('rm -rf .smake')
else:
    config.run(args.target, args.mode, args.verbose)