#!/usr/bin/python
# the cook build system 0.1.0
#
# Copyright 2023 Alejandro W. Sior
# 
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject
# to the following conditions:

# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR
# ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF
# CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import subprocess
import sys
import os
from pathlib import Path
import shutil
import argparse

def error(*args):
    print(*args, file = sys.stderr)

class ProgramNotFoundException(Exception):
    pass

class AttrDict(dict):
    def __init__(self, *args, **kwargs):
        super(AttrDict, self).__init__(*args, **kwargs)
        self.__dict__ = self

def find_program(prog, required=True):
    path = shutil.which(prog)
    if path == None and required:
        raise ProgramNotFoundException(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)
DEFMKDESC = lambda o, i, cmd: "%s" % (cmd)
DEFMKOUT = lambda i: "%s" % (i)

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, depends = [], extra_args = []):
        super().__init__(path)

        # The generator that created this process
        self.processor = processor

        # The inputs
        self.inputs = File.resolve(inputs)

        # Implicit dependencies
        self.depends = File.resolve(depends)
        if processor.exe_is_process:
            self.depends += [processor.exe]

        # The extra arguments taken
        self.extra_args = flatten(extra_args)

    def gen_makefile(self):
        inputs = " ".join([str(x) for x in self.inputs])
        depends = " ".join([str(x) for x in self.depends])
        extra_args = " ".join([x for x in self.extra_args])

        out = "%s: %s %s\n" % (self.path, inputs, depends)
        out += "\t@$(%s) %s\n" % (self.processor.name, self.processor.mkargs("$@", inputs, extra_args))
        out += "\t@echo '%s'" % (self.processor.mkdesc("$@", inputs, extra_args))

        return out

    def gen_ninja(self):
        inputs = " ".join([str(x) for x in self.inputs])
        depends = ""
        if self.depends != []:
            depends = "| " + " ".join([str(x) for x in self.depends])
        extra_args = " ".join(self.extra_args)

        out = "build %s: %s %s %s" % (self.path, self.processor.name, inputs, depends)
        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, mkout=DEFMKOUT, mkdesc=DEFMKDESC):
        # The name of the processor
        self.name = name

        # Executable used in the process
        self.exe = File(exe)
        if issubclass(type(self.exe), Target):
            self.exe_is_process = True
        else:
            self.exe_is_process = False

        # The arguments that is used to launch the process
        self.mkargs = wrap_lambda(mkargs)

        # The way to generate output path
        self.mkout = wrap_lambda(mkout)

        # The description
        self.mkdesc = wrap_lambda(mkdesc)

    def __call__(self, *args, **kwargs):
        return self.gen(*args, **kwargs)

    def gen(self, name, *args, **kwargs):
        inputs = File.resolve(args)
        path = BUILD / Path(".").resolve().relative_to(ROOT) / self.mkout(name)

        return Target(path, self, inputs, **kwargs)

    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.mkdesc("$out", "$in", cmd))
        return out

class Vec(Processor):
    def __init__(self, *args, mkout, **kwargs):
        super().__init__(*args, mkout = mkout, **kwargs)

    def __call__(self, *args, **kwargs):
        return self.gen(*args, **kwargs)

    def gen(self, *args, **kwargs):
        args = File.resolve(args)

        targets = []
        for i in args:
            path = BUILD / i.path.parent / self.mkout(i.path)

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

        return targets

class CTarget(Target):
    def __init__(self, *args, header_depends = [], **kwargs):
        super().__init__(*args, **kwargs)
        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, *args, mkinc, mkdep=None, depstyle=None, **kwargs):
        super().__init__(*args, **kwargs)
        self.mkinc = mkinc
        self.mkdep = mkdep
        self.depstyle = depstyle

    def __call__(self, *args, **kwargs):
        return self.gen(*args, **kwargs)

    def gen(self, *args, extra_args = [], include_dirs = [], dependencies = [], **kwargs):
        inputs = File.resolve(args)
        dependencies = flatten(dependencies)

        include_dirs += [d.include_dirs for d in dependencies if d.include_dirs]
        include_dirs = File.resolve(include_dirs)
        include_args = [self.mkinc(str(x)) for x in include_dirs]
        extra_args += [d.cc_args for d in dependencies if d.cc_args]
        extra_args += include_args

        targets = []
        for input in inputs:
            out = self.mkout(input.path)
            path = BUILD / out

            header_depends = []
            if self.mkdep:
                header_depends += [BUILD / self.mkdep(out)]

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

        return targets

    def gen_ninja(self):
        out = super().gen_ninja()

        if self.depstyle == "gcc":
            out += "  deps = gcc\n"
        elif self.depstyle == "msvc":
            out += "  deps = msvc\n"
        else:
            raise Exception("ninja dependencies for selected depstyle is unimplemented")

        if self.mkdep and self.depstyle == "gcc":
            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, *args, mklib, **kwargs):
        super().__init__(*args, **kwargs)

        self.mklib = mklib

    def __call__(self, *args, **kwargs):
        return self.gen(*args, **kwargs)

    def gen(self, name, *args, libs=[], dependencies = [], extra_args = [], **kwargs):
        libs = flatten(libs)
        dependencies = flatten(dependencies)
        
        libs += [d.libs for d in dependencies if d.libs]
        extra_args += [d.ld_args for d in dependencies if d.ld_args]
    
        # 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, **kwargs)

def CCStyleToolchain(name, gccname, arname, suffix="", c_args=[], ld_args=[]):
    ccname = name + suffix
    exename = name + "exe" + suffix
    libname = name + "lib" + suffix
    shlibname = name + "shlib" + suffix

    c_args = " ".join(flatten(c_args))
    ld_args = " ".join(flatten(ld_args))

    gcc = find_program(gccname)
    ar = find_program(arname)

    cc = CC(
        name=ccname,
        exe=gcc,
        mkargs=lambda o, i, ea: "-MD -MF %s.d -o %s -c %s %s %s" % (o, o, i, c_args, ea),
        mkout=lambda i: "%s.o" % (i),
        mkinc=lambda inc: "-I%s" % (inc),
        mkdep=lambda i: "%s.d" % (i),
        depstyle="gcc",
        mkdesc=lambda o, i, cmd: "CC %s" % (o),
    )
    exe = CExe(
        name=exename,
        exe=gcc,
        mkargs=lambda o, i, ea: "-o %s -Wl,--start-group %s -Wl,--end-group %s %s" % (o, i, ld_args, ea),
        mklib=lambda l: "-l" + l,
        mkdesc=lambda o, i, cmd: "LD %s" % (o)
    )
    lib = CExe(
        name=libname,
        exe=ar,
        mkargs=lambda o, i, ea: "-rcs %s %s %s" % (o, i, ea),
        mklib=lambda l: "", # TODO: this should not be needed
        mkdesc=lambda o, i, cmd: "AR %s" % (o),
        mkout=lambda n: "lib%s.a" % (n)
    )
    shlib = CExe(
        name=shlibname,
        exe=gcc,
        mkargs=lambda o, i, ea: "-o %s -Wl,--start-group %s -Wl,--end-group %s -fPIC --shared %s" % (o, i, ld_args, ea),
        mklib=lambda l: "-l" + l,
        mkdesc=lambda o, i, cmd: "SO %s" % (o),
        mkout=lambda n: "lib%s.so" % (n)
    )
    return AttrDict({'cc': cc, 'exe': exe, 'lib': lib, 'shlib': shlib})

def GccToolchain(*args, **kwargs):
    return CCStyleToolchain("gcc", "gcc", "gcc-ar", *args, **kwargs)

def ClangToolchain(*args, **kwargs):
    return CCStyleToolchain("clang", "clang", "llvm-ar", *args, **kwargs)

def CCToolchain(*args, **kwargs):
    return CCStyleToolchain("cc", "cc", "ar", *args, **kwargs)

def MSVCStyleToolchain(name, gccname, arname, ldname, suffix="", c_args=[], ld_args=[]):
    ccname = name + suffix
    exename = name + "exe" + suffix
    libname = name + "lib" + suffix
    shlibname = name + "shlib" + suffix

    c_args = " ".join(flatten(c_args))
    ld_args = " ".join(flatten(ld_args))

    cl = find_program(gccname)
    lib = find_program(arname)
    link = find_program(ldname)

    cc = CC(
        name=ccname,
        exe=cl,
        mkargs=lambda o, i, ea: "/showIncludes /nologo /D_USRDLL /D_WINDLL /Fo%s /c %s /MT %s %s" % (o, i, c_args, ea),
        mkout=lambda i: "%s.obj" % (i),
        mkinc=lambda inc: "/I%s" % (inc),
        mkdesc=lambda o, i, cmd: "CC %s" % (o),
        depstyle="msvc"
    )
    exe = CExe(
        name=exename,
        exe=link,
        mkargs=lambda o, i, ea: "/NOLOGO /OUT:%s %s %s %s" % (o, i, ld_args, ea),
        mklib=lambda l: l,
        mkdesc=lambda o, i, cmd: "LD %s" % (o),
        mkout=lambda n: "%s.exe" % (n),
    )
    lib = CExe(
        name=libname,
        exe=lib,
        mkargs=lambda o, i, ea: "/NOLOGO /OUT:%s %s %s" % (o, i, ea),
        mklib=lambda l: "", # TODO: this should not be needed
        mkdesc=lambda o, i, cmd: "AR %s" % (o),
        mkout=lambda n: "lib%s.lib" % (n)
    )
    shlib = CExe(
        name=shlibname,
        exe=cl,
        mkargs=lambda o, i, ea: "/nologo /D_USRDLL /D_WINDLL %s /MT /link /DLL /OUT:%s.dll /IMPLIB:%s %s %s" % (i, o, o, ld_args, ea),
        mklib=lambda l: "", # TODO: make this not needed, or perhaps yes by adding .lib but this seems weird
        mkdesc=lambda o, i, cmd: "SO %s" % (o),
        mkout=lambda n: "%s.lib" % (n)
    )
    return AttrDict({'cc': cc, 'exe': exe, 'lib': lib, 'shlib': shlib})

def MSVCToolchain(*args, **kwargs):
    return MSVCStyleToolchain("msvc", "cl", "lib", "link", *args, **kwargs)
    
def ClangClToolchain(*args, **kwargs):
    # This toolchain depends on lib.exe shipped by microsoft. Is there a llvm way?
    return MSVCStyleToolchain("clang-cl", "clang-cl", "lib", "lld-link", *args, **kwargs)

def CppToolchain(*args, **kwargs):
    return CCStyleToolchain("cpp", "c++", "ar", *args, **kwargs)

def GppToolchain(*args, **kwargs):
    return CCStyleToolchain("gpp", "g++", "gcc-ar", *args, **kwargs)

def ClangppToolchain(*args, **kwargs):
    return CCStyleToolchain("clangpp", "clang++", "llvm-ar", *args, **kwargs)


CTOOLCHAINS = {
    'msvc': MSVCToolchain,
    'clang-cl': ClangClToolchain,
    'cc': CCToolchain,
    'clang': ClangToolchain,
    'gcc': GccToolchain
}

CXXTOOLCHAINS = {
    'msvc': MSVCToolchain,
    'clang-cl': ClangClToolchain,
    'c++': CppToolchain,
    'g++': GppToolchain,
    'clang++': ClangppToolchain
}


def CToolchain(name="system", suffix="", c_args=[], ld_args=[], tdb=CTOOLCHAINS):
    if suffix == "" and (c_args != [] or ld_args != []):
        error("warning: instanciating default toolchain %s with non-default arguments" % (name))
        error("note: to fix this, add a suffix")

    if name != "system":
        if not name in tdb:
            raise Exception("No such C toolchain available")

        return tdb[name](suffix, c_args, ld_args)

    for t in tdb.values():
        try:
            return t(suffix, c_args, ld_args)
        except ProgramNotFoundException as e:
            continue
        except Exception as e:
            raise e

    raise Exception("No C toolchain found on system")

def CppToolchain(*args, tdb=CXXTOOLCHAINS, **kwargs):
    return CToolchain(*args, tdb = tdb, **kwargs)

class CMod:
    def __init__(self, parent, name):
        self.path = parent.path / name if parent else ROOT
    
        self.name = name
        self.parent = parent
        if self.parent:
            self.parent.children[name] = self
        self.children = dict()
    
    def __call__(self, path):
        components = path.split("/")
        
        resolved = self
        
        for i, component in enumerate(components):
            if i == 0 and component == "":
                resolved = _root
            elif component == "" or component == ".":
                resolved = resolved
            elif component == "..":
                resolved = resolved.parent
            elif component in resolved.children:
                resolved = resolved.children[component]
            else:
                m = CMod(resolved, component)
                m.dive()
                resolved = m
        
        return resolved

    def __getattr__(self, attr):
        if self.parent:
            return getattr(self.parent, attr)
        else:
            return None
    
    def dive(self):
        global this
        # Save the previous location
        PWD = Path.cwd()
        previous_this = this
        this = self
    
        # Go into subdir
        sub = self.path
        os.chdir(sub)
    
        exec(open('.build').read(), globals())
    
        # Revert the current directory;
        os.chdir(PWD)
        this = previous_this
        
global _root
_root = CMod(None, "")

global this
this = _root

global cook
cook = AttrDict({
    'c': CToolchain(),
    'cpp': CppToolchain()
})

# Set default toolchains
this.c = cook.c
this.cpp = cook.cpp

pkgconfig = find_program("pkg-config", False)

def subproc(*args, **kwargs):
    return subprocess.check_output(*args, **kwargs, stderr=subprocess.DEVNULL)

# XXX: make name a variadic
def CDependency(name):
    if type(name) == list or type(name) is tuple:
        return [CDependency(n) for n in name]

    module = CMod(None, name)

    if pkgconfig:
        try:
            inc = subproc([pkgconfig.path, "--cflags", name]).strip().decode("utf-8")
            lib = subproc([pkgconfig.path, "--libs", name]).strip().decode("utf-8")
            
            module.cc_args = inc if inc != "" else []
            module.ld_args = lib if lib != "" else []
            
            return module
        except subprocess.CalledProcessError:
            pass
    
    #if prog:
    # XXX: pkg-config
    # XXX: windows stuff

    # Use the system way of getting to know libraries (no includes specified)
    module.libs = name
    return module
    

# cc, exe, lib, shlib = CToolchain()

# #global root, mod
# root = CModule("root", cc, exe, lib, shlib)
# mod = root

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

def subdir(dir):
    # Save the previous location
    PWD = Path.cwd()

    # Go into subdir
    sub = PWD / dir
    os.chdir(sub)

    # Set helper path
    global CWD
    CWD = Path.cwd().relative_to(ROOT)
    exec(open('.build').read(), globals())

    # Revert the current directory;
    os.chdir(PWD)

this.dive()

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 0.1.0")
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))