#!/usr/bin/python

import os
from pathlib import Path
from threading import local
import shutil
import argparse

def find_program(prog, required=True):
    path = shutil.which(prog)
    if path == None and required:
        raise Exception("Required program '%s' not found on the system" % (prog))
    return File(path)

def flatten(el):
    if type(el) == list or type(el) == tuple:
        return [a for b in el for a in flatten(b)]
    else:
        return [el]

ROOT = Path.cwd()
BUILD = Path("bld").resolve().relative_to(ROOT)
DEFDESC = lambda o, i, cmd: "%s" % (cmd)

#self, name, exe, mkargs, mkout, mkinc, desc=DEFDESC

C_COMPILERS = {
    'gcc': lambda: CC(
        "gcc",
        exe=find_program("gcc"),
        mkargs=lambda o, i, ea: "-MD -MF %s.d -o %s -c %s %s" % (o, o, i, ea),
        mkout=lambda i: "%s.o" % (i),
        mkinc=lambda i: "-I" + i,
        mkdep=lambda o: "%s.d" % (o),
        depstyle="gcc",
        desc=lambda o, i, cmd: "CC %s" % (o),
    ),
    'clang': lambda: CC(
        "clang",
        find_program("clang"),
        lambda o, i, ea: "-MD -MF %s.d -o %s -c %s %s" % (o, i, ea),
        mkargs=lambda o, i, ea: "-MD -MF %s.d -o %s -c %s %s" % (o, o, i, ea),
        mkout=lambda i: i.name + ".o",
        mkinc=lambda i: "-I" + i,
        mkdep=lambda o: "%s.d" % (o),
        depstyle="gcc",
        desc=lambda o, i, cmd: "CC %s" % (o),
    )
}

C_LINKERS = {
    'gcc': lambda: CExe(
        "ld-gcc",
        find_program("gcc"),
        mkargs=lambda o, i, ea: "-o %s %s %s" % (o, i, ea),
        mklib=lambda l: "-l" + l,
        desc=lambda o, i, cmd: "LD %s" % (o)
    ),
    'clang': lambda: CExe(
        "ld-clang",
        find_program("clang"),
        mkargs=lambda o, i, ea: "-o %s %s %s" % (o, i, ea),
        mklib=lambda l: "-l" + l,
        desc=lambda o, i, cmd: "LD %s" % (o)
    )
}

def wrap_lambda(s):
    if callable(s):
        return s
    else:
        return lambda *args: s


class FileMeta(type):
    def __call__(cls, path, *rest, **krest):
        # If none, return none
        if path == None:
            return None
    
        # If already a file, return it
        if issubclass(type(path), File):
            return path
    
        # Adjust name to be a path
        if type(path) == str:
            p = Path(path).resolve()
            if p.is_relative_to(ROOT):
                path = p.relative_to(ROOT)
            else:
                path = p
    
        idx = str(path) + cls.__name__
    
        if idx in cls.fdb:
            return cls.fdb[idx]
        
        obj = cls.__new__(cls, path, *rest, **krest)
        cls.__init__(obj, path, *rest, **krest)
        cls.fdb[idx] = obj
        return obj

class File(metaclass = FileMeta):
    fdb = dict()

    def __init__(self, path):
        # The name corresponding to
        # the target: this can be
        # a path or an abstract name
        self.path = path
                
    def __str__(self):
        return str(self.path)

    def resolve(*args):
        return [File(file) for file in flatten(args)]

class Target(File):
    def __init__(self, path, processor, inputs, extra_args = []):
        super().__init__(path)
        
        # The generator that created this process
        self.processor = processor
        
        # The inputs
        self.inputs = File.resolve(inputs)
        
        # The extra arguments taken
        self.extra_args = extra_args
    
    def gen_makefile(self):
        inputs = " ".join([str(x) for x in self.inputs])
        extra_args = " ".join([x for x in self.extra_args])

        out = "%s: %s\n" % (self.path, inputs)
        out += "\t@$(%s) %s\n" % (self.processor.name, self.processor.mkargs("$@", inputs, extra_args))
        out += "\t@echo '%s'" % (self.processor.desc("$@", inputs, extra_args))
        
        return out
    
    def gen_ninja(self):
        inputs = " ".join([str(x) for x in self.inputs])
        extra_args = " ".join(self.extra_args)
        
        out = "build %s: %s %s" % (self.path, self.processor.name, inputs)
        if self.extra_args != []:
            out += "\n  extra = %s" % (extra_args)

        return out

class ProcessorMeta(type):
    def __call__(cls, name, *rest, **krest):
        # If already a file, return it
        if issubclass(type(name), Processor):
            return name
    
        if name in cls.mdb:
            return cls.mdb[name]
        
        obj = cls.__new__(cls, name, *rest, **krest)
        cls.__init__(obj, name, *rest, **krest)
        cls.mdb[name] = obj
        return obj

class Processor(metaclass = ProcessorMeta):
    mdb = dict()

    def __init__(self, name, exe, mkargs, desc=DEFDESC):
        # The name of the processor
        self.name = name
        
        # Executable used in the process
        self.exe = File(exe)
        
        # The arguments that is used to launch the process
        self.mkargs = wrap_lambda(mkargs)
        
        # The description
        self.desc = wrap_lambda(desc)

    def __call__(self, name, *args, extra_args = []):
        return self.gen(name, *args, extra_args = extra_args)
    
    def gen(self, name, *args, extra_args = []):
        inputs = File.resolve(args)
        path = BUILD / Path.cwd().relative_to(ROOT) / name
        
        return Target(path, self, inputs, extra_args)

    def gen_makefile(self):
        return "%s = %s" % (self.name, self.exe)

    def gen_ninja(self):
        out = "rule %s\n" % (self.name)
        
        cmd = "%s %s" % (self.exe, self.mkargs("$out", "$in", "$extra"))
        
        out += "  command = %s\n" % (cmd)
        out += "  description = %s\n" % (self.desc("$out", "$in", cmd))
        return out

class Vec(Processor):
    def __init__(self, name, exe, mkargs, mkout, desc=DEFDESC):
        super().__init__(name, exe, mkargs, desc=desc)
        self.mkout = wrap_lambda(mkout)
    
    def __call__(self, *args, extra_args = []):
        return self.gen(*args, extra_args = extra_args)
    
    def gen(self, *args, extra_args = []):
        args = File.resolve(args)
        extra_args = flatten(extra_args)
        
        targets = []
        for i in args:
            path = BUILD / i.path.parent / self.mkout(i.path)

            targets.append(Target(path, self, [i], extra_args))

        return targets

class CTarget(Target):
    def __init__(self, path, processor, inputs, header_depends, extra_args = []):
        super().__init__(path, processor, inputs, extra_args = extra_args)
        self.header_depends = File.resolve(header_depends)
    
    def gen_makefile(self):
        out = []

        if self.header_depends != []:
            out += ["-include %s\n" % hdep for hdep in self.header_depends]
        
        out += super().gen_makefile()
        return "".join(out)

class CC(Vec):
    def __init__(self, name, exe, mkargs, mkout, mkinc, mkdep=None, depstyle=None, desc=DEFDESC):
        super().__init__(name, exe, mkargs, mkout, desc=desc)
        self.mkinc = mkinc
        self.mkdep = mkdep
        self.depstyle = depstyle
    
    def __call__(self, *args, extra_args = [], include_dirs = []):
        return self.gen(*args, extra_args = extra_args, include_dirs = include_dirs)
    
    def gen(self, *args, extra_args = [], include_dirs = []):
        inputs = File.resolve(args)
    
        extra_args = flatten(extra_args)
        include_dirs = File.resolve(include_dirs)

        include_args = [self.mkinc(str(x)) for x in include_dirs]
        extra_args += include_args
        
        targets = []
        for input in inputs:
            out = self.mkout(input.path)
            path = BUILD / out
            hdep = BUILD / self.mkdep(out)

            targets.append(CTarget(path, self, input, header_depends = hdep, extra_args = extra_args))

        return targets
        
    def gen_ninja(self):
        out = super().gen_ninja()
        
        if self.mkdep and self.depstyle == "gcc":
            out += "  deps = gcc\n"
            out += "  depfile = %s\n" % (self.mkdep("$out"))
        elif self.mkdep:
            raise Exception("ninja dependencies for selected depstyle is unimplemented")
        
        return out

class CExe(Processor):
    def __init__(self, name, exe, mkargs, mklib, desc=DEFDESC):
        super().__init__(name, exe, mkargs, desc = desc)
        
        self.mklib = mklib
    
    def __call__(self, name, *args, libs = [], extra_args = []):
        return self.gen(name, *args, libs = libs, extra_args = extra_args)
    
    def gen(self, name, *args, libs, extra_args = []):
        extra_args = flatten(extra_args)
        libs = flatten(libs)
        
        # Note: this is bad, perhaps take a "depends" also/instead
        # XXX Handling of libraries is temporary
        libs_args = [self.mklib(lib) for lib in libs]
        extra_args += libs_args
        
        return super().gen(name, *args, extra_args = extra_args)

def find_one(names):
    for name in names:
        program = find_program(name, required=False)
        if program:
            return name
    return None

class System:
    def c_compiler(c_compilers=C_COMPILERS.keys()):
        c_compiler = find_one(c_compilers)
        return C_COMPILERS[c_compiler]()

    def c_linker(c_linkers=C_LINKERS.keys()):
        c_linker = find_one(c_linkers)
        return C_LINKERS[c_linker]()

def subdir(dir):
    # Save the previous location
    PWD = Path.cwd()
    
    # Go into subdir
    sub = PWD / dir
    os.chdir(sub)
    
    # Set helper path
    CWD = Path.cwd().relative_to(ROOT)
    exec(open('.build').read())
    
    # Revert the current directory;
    os.chdir(PWD)

subdir(".")

def gen_bld():
    return [target.path.parent.mkdir(parents=True, exist_ok=True)
        for target in File.fdb.values()
        if issubclass(type(target), Target)]

def gen_makefile():
    out = [mach.gen_makefile() + "\n" for mach in Processor.mdb.values()]
    out += ["\n\n" + target.gen_makefile() for target in File.fdb.values() if issubclass(type(target), Target)]
    
    return "".join(out)
    
def gen_ninja():
    out = [mach.gen_ninja() + "\n" for mach in Processor.mdb.values()]
    out += [target.gen_ninja() + "\n" for target in File.fdb.values() if issubclass(type(target), Target)]
    
    return "".join(out)

parser = argparse.ArgumentParser(description="cook build system")
parser.add_argument("-G", type=str, metavar="backend", help="Backend to generate for: ninja or makefile")
args = parser.parse_args()

backend = "ninja"
if args.G:
    backend = args.G

gen_bld()
open(ROOT / BUILD / ".gitignore", "w").write("*")
if backend == "ninja":
    open(ROOT / 'build.ninja', 'w').write(gen_ninja())
elif backend == "makefile":
    open(ROOT / 'makefile', 'w').write(gen_makefile())
else:
    print("No such backend %s" % (backend))