#!/usr/bin/env python3

import argparse
import collections
import dataclasses
import glob
import itertools
import os
import re


class Config:
    def __init__(self, builddir, outdir):
        self.builddir = builddir
        self.outdir = outdir

        if not os.path.exists(builddir):
            os.mkdir(builddir)
        if not os.path.exists(outdir):
            os.mkdir(outdir)

    toolchain_path = f"/opt/xpack-riscv-none-elf-gcc-14.2.0-3"
    toolchain_prefix = "riscv-none-elf-"
    hostcflags = "-g -std=c++20 -fprofile-instr-generate -fcoverage-mapping"
    hostlibs = "-lgtest -lgmock -lgtest_main"
    hostldflags = "-fprofile-instr-generate -fcoverage-mapping"
    include_dirs = [
        "hal",
        "hal/uart",
        "hal/lib/common",
    ]
    common_flags = [
        "-g",
        "-Wall",
        "-Wextra",
        "-flto",
        "-march=rv32i_zicsr",
        "-ffunction-sections",
        "-Oz",
    ]
    ldflags = [
        "-Oz",
        "-g",
        "-Wl,--gc-sections",
        "-Wl,--print-memory-usage",
        "-specs=nano.specs",
        "-flto",
        "-march=rv32i_zicsr",
    ]
    project_flags = [
    ]

    @property
    def include_flags(self):
        return [f"-I{os.path.relpath(i, self.builddir)}" for i in self.include_dirs]

    @property
    def cpp_flags(self):
        return ["-DNDEBUG"] + self.include_flags + self.project_flags

    @property
    def cc_flags(self):
        return self.common_flags + self.cpp_flags

    @property
    def cxx_flags(self):
        return self.common_flags + [
                "-std=c++20",
                "-fno-rtti",
                "-fno-exceptions",
                "-Wno-missing-field-initializers",
                ] + self.cpp_flags

    def relbuild(self, path):
        return os.path.relpath(path, self.builddir)

def gen_rules(config):
    tools = {"cxx": "g++", "cc": "gcc", "as": "as", "objcopy": "objcopy"}

    tc_path = config.toolchain_path
    tc_prefix = config.toolchain_prefix

    rules = f"""
rule cxx
  command = $cxx -MMD -MT $out -MF $out.d {' '.join(config.cxx_flags)} -c $in -o $out
  description = CXX $out
  depfile = $out.d
  deps = gcc

rule cc
  command = $cc -MMD -MT $out -MF $out.d {' '.join(config.cc_flags)} -c $in -o $out
  description = CC $out
  depfile = $out.d
  deps = gcc

rule as
  command = $as $in -o $out
  description = AS $out

rule link
  command = $cxx {' '.join(config.ldflags)} -Wl,-T$linker_script -o $out $in $libs
  description = LINK $out

rule objcopy
  command = $objcopy -O binary $in $out
  description = OBJCOPY $out

rule hostcxx
  command = clang++ -MMD -MT $out -MF $out.d {config.hostcflags} -c $in -o $out
  description = HOSTCXX $out
  depfile = $out.d
  deps = gcc

rule hostlink
  command = clang++ {config.hostldflags} -o $out $in {config.hostlibs}
  description = HOSTLINK $out

rule profdata
  command = llvm-profdata merge -sparse $profraw -o $out
  description = PROFDATA

rule cov
  command = llvm-cov show --output-dir cov -format html --instr-profile $profdata $objects && touch $out
  description = COV

rule hosttest
  command = LLVM_PROFILE_FILE=$in.profraw ./$in && touch $out
"""

    for var, tool in tools.items():
        toolpath = os.path.join(tc_path, "bin", f"{tc_prefix}{tool}")
        yield f"{var} = {toolpath}"

    for line in rules.splitlines():
        yield line


def get_suffix_rule(filename, cxx_rule="cxx", cc_rule="cc"):
    suffix = filename.split(".")[-1]
    return collections.defaultdict(
        lambda: None,
        {
            "c": cc_rule,
            "cc": cxx_rule,
            "cpp": cxx_rule,
            "s": "as",
        },
    )[suffix]


def make_cxx_rule(name, cflags=()):
    cflags = " ".join(cflags)
    rule = f"""
rule {name}
  command = $cxx -MMD -MT $out -MF $out.d {cflags} -c $in -o $out
  description = CXX $out
  depfile = $out.d
  deps = gcc
"""
    return rule.splitlines()


def make_cc_rule(name, cflags=()):
    cflags = " ".join(cflags)
    rule = f"""
rule {name}
  command = $cc -MMD -MT $out -MF $out.d {cflags} -c $in -o $out
  description = CC $out
  depfile = $out.d
  deps = gcc
"""
    return rule.splitlines()


@dataclasses.dataclass
class SourceSet:
    name: str
    sources: list[str]
    cflags: list[str]

    def get_objects(self, config):
        for s in self.sources:
            yield (config.relbuild(s), re.sub(r"\.\w+", ".o", s))

    def get_targets(self, config):
        return list(zip(*self.get_objects(config)))[1]


def source_set(name, sources, cflags=()):
    return SourceSet(name, sources, cflags)


def build_source_set(source_set):
    def __f(config):
        cxx_rule = "cxx"
        cc_rule = "cc"
        lines = []
        if source_set.cflags:
            cxx_rule = f"cxx_{name}"
            lines += make_cxx_rule(cxx_rule, cflags=cflags + config.cxx_flags)
            cc_rule = f"cc_{name}"
            lines += make_cc_rule(cc_rule, cflags=cflags + config.cc_flags)
        for line in lines:
            yield line

        for i, o in source_set.get_objects(config):
            rule = get_suffix_rule(i, cxx_rule=cxx_rule, cc_rule=cc_rule)
            if rule is None:
                continue
            yield f"build {o}: {rule} {i}"

    return __f


def build_image(source_set, elf_out, linker_script, dependencies=(), bin_out=None):
    def __f(config):
        # to make it builddir-relative
        lscript = config.relbuild(linker_script)
        elfout = config.relbuild(os.path.join(config.outdir, elf_out))

        for l in build_source_set(source_set)(config):
            yield l

        objects = source_set.get_targets(config)
        for dep in dependencies:
            objects += dep.get_targets(config)
        objects = " ".join(objects)

        yield f"build {elfout}: link {objects} | {lscript}"
        yield f"  linker_script = {lscript}"
        if bin_out is not None:
            binout = config.relbuild(os.path.join(config.outdir, bin_out))
            yield f"build {binout}: objcopy {elfout}"

    return __f


def build_test(name, sources):
    builds = [
        (os.path.relpath(s, builddir), f"{name}_" + re.sub(r"\.\w+", ".o", s))
        for s in sources
    ]

    out = name

    for i, o in builds:
        rule = "hostcxx"
        yield f"build {o}: {rule} {i}"

    objects = " ".join(b[1] for b in builds)

    yield f"build {out}: hostlink {objects}"
    yield f"build {out}.run: hosttest {out}"


def make_coverage(binaries):
    bins = " ".join(binaries)
    profraw = " ".join(f"{x}.profraw" for x in binaries)
    objects = " ".join(f"--object {x}" for x in binaries)
    testruns = " ".join(f"{x}.run" for x in binaries)
    yield f"build profdata: profdata | {testruns}"
    yield f"  profraw = {profraw}"
    yield f"build cov/index.html: cov {bins} | profdata"
    yield f"  profdata = profdata"
    yield f"  objects = {objects}"


hal = source_set("hal", [
        "hal/intc.cc",
        "hal/interrupts.cc",
        "hal/start.cc",
        "hal/lib/common/xil_assert.c",
        "hal/uart/xuartlite.c",
        "hal/uart/xuartlite_stats.c",
        "hal/uart/xuartlite_intr.c",
        ])

bootloader = source_set("bootloader", glob.glob("./bootloader/**/*.cc", recursive=True))
bootloader_image = build_image(
    bootloader,
    dependencies=[hal],
    elf_out="bootloader.elf",
    linker_script="bootloader/bootloader.ld",
)

def app_image(app, sources=None):
    if sources is None:
        sources = glob.glob(f"./apps/{app}/**/*.cc", recursive=True)
    return build_image(
            source_set(app, sources),
            linker_script="apps/app.ld",
            dependencies=[hal],
            elf_out=f"{app}.elf",
            bin_out=f"{app}.bin",
    )

all = [
        build_source_set(hal),
        bootloader_image,

        app_image("helloworld"),
        app_image("timer"),
        app_image("uart"),
        app_image("async", sources=[
            "apps/async/async.cc",
            "apps/async/lock.cc",
            "apps/async/main.cc",
            "apps/async/trace.cc",
            "apps/async/uart.cc",
            ]),
]


def parse_args():
    parser = argparse.ArgumentParser(description='Generate ninja build files.')
    parser.add_argument('--version', required=True,
                        help='version tag (typically from `git describe`)')
    parser.add_argument('--build-dir', default='build', help='build directory')
    parser.add_argument('--out-dir', default='out', help='output directory')
    return parser.parse_args()


def main():
    args = parse_args()
    config = Config(builddir=args.build_dir, outdir=args.out_dir)
    config.project_flags.append(f'-DGIT_VERSION_TAG=\\"{args.version}\\"')
    header = gen_rules(config)
    lines = itertools.chain(header, *(f(config) for f in all))

    with open(os.path.join(config.builddir, "build.ninja"), "w") as f:
        f.write("\n".join(lines))
        f.write("\n")

    print(
        f'Configure done. Build with "ninja -C {config.builddir}". Output will be in {config.outdir}/'
    )


if __name__ == "__main__":
    main()