#!/usr/bin/env python

import os.path
from pathlib import Path
from invoke import task
from textwrap import dedent
import shutil

def ensure_build_dirs(root : Path, *subdirs : list[Path]):
    if root.is_file():
        root = root.parent / root.stem
    ret = []
    for d in subdirs:
        out = "out" / root / d
        out.mkdir(mode=0o750, parents=True, exist_ok=True)
        ret.append(out)
    return ret

def bsc_root(c):
    out = c.run("bsc --help", hide="both")
    for l in out.stdout.splitlines():
        dir_prefix = "Bluespec directory: "
        if l.startswith(dir_prefix):
            return Path(l[len(dir_prefix):])
    raise RuntimeError("Couldn't locate Bluespec root dir")

def find_verilog_modules(c, modules):
    libpaths = [Path("lib"), bsc_root(c) / "Verilog"]
    ret = []
    for module in modules:
        module_path = None
        for p in libpaths:
            f = p / Path(module).with_suffix(".v")
            if f.is_file():
                module_path = f
                break
        if module_path is None:
            raise RuntimeError(f"Cannot find verilog module {module} in {libpaths}")
        ret.append(module_path)
    return ret

def print_filtered_paragraphs(s, *paragraph_prefixes, common_prefix=""):
    for block in s.split("\n\n"):
        for p in paragraph_prefixes:
            if block.removeprefix(common_prefix).startswith(p):
                print('\n'.join(s.removeprefix(common_prefix) for s in block.splitlines()))
                print("")

def expand_build_target(target):
    t = Path(target)
    if t.is_dir():
        yield from t.glob("**/*.bsv")
    elif t.is_file() and t.suffix == ".bsv":
        yield t
    else:
        raise ValueError(f"Unknown target type {t}")

def resolve_synth_target(target):
    target = Path(target)
    if '/' not in str(target):
        target = "hardware" / target
    if target.is_dir():
        target /= "Top.bsv"
    if not target.is_file():
        raise ValueError(f"Unknown target type {target}")
    return target

def expand_test_target(target):
    t = Path(target)
    if t.is_dir():
        yield from t.glob("**/*_Test.bsv")
    elif t.is_file() and t.name.endswith("_Test.bsv"):
        yield t
    else:
        raise ValueError(f"Unknown target type {t}")

def phase(name):
    print("")
    print("*"*(len(name)+6))
    print(f"** {name} **")
    print("*"*(len(name)+6))
    print("")

@task
def build(c, target="."):
    phase("Compile Bluespec")

    verilog_files = []
    for target in expand_build_target(target):
        out_info, out_verilog, out_bsc = ensure_build_dirs(target, "info", "verilog", "bsc")
        print(f"Building {target}")
        c.run(f"bsc -aggressive-conditions -check-assert -remove-dollar -remove-empty-rules -remove-false-rules -remove-starved-rules -remove-unused-modules -show-method-conf -show-method-bvi -u -verilog -info-dir {out_info} -vdir {out_verilog} -bdir {out_bsc} -p {target.parent}:lib:%/Libraries -show-module-use -show-compiles {target}")

        module_name = Path(f"mk{target.stem}")
        verilog_main_file = out_verilog / module_name.with_suffix(".v")
        if verilog_main_file.is_file():
            verilog_files.append(verilog_main_file)
        use_file = out_verilog / module_name.with_suffix(".use")
        if use_file.is_file():
            with open(out_verilog / module_name.with_suffix(".use")) as f:
                verilog_files.extend(find_verilog_modules(c, f.read().splitlines()))

    if verilog_files:
        print("\nVerilog files for synthesis:")
        for v in verilog_files:
            print("  "+str(v))

    return verilog_files

@task
def synth(c, target):
    target = resolve_synth_target(target)

    out_info, out_verilog, out_bsc, out_yosys, out_nextpnr = ensure_build_dirs(target, "info", "verilog", "bsc", "yosys", "nextpnr")

    module_name = Path(f"mk{target.stem}")
    verilog_files = build(c, target)
    if not verilog_files:
        print("\nWARNING: bluespec compile didn't output any verilog files, did you (* synthesize *) something?")
        return

    phase("Logic synthesis")
    yosys_script = out_yosys / "script.ys"
    yosys_json = out_yosys / module_name.with_suffix(".json")
    yosys_report = out_yosys / module_name.with_suffix(".stats")
    yosys_log = out_yosys / module_name.with_suffix(".log")
    yosys_preprocessed = out_yosys / module_name.with_suffix(".v")
    yosys_preprocessed_rtl = out_yosys / module_name.with_suffix(".rtlil")
    yosys_preprocessed_graph = out_yosys / module_name.stem
    yosys_compiled = out_yosys / f"{module_name.stem}_compiled.v"
    with open(yosys_script, "w", encoding="UTF-8") as f:
        f.write(dedent(f"""\
            read_verilog -sv -defer {' '.join(str(f) for f in verilog_files)}
            hierarchy -top {module_name}
            design -push-copy
            synth_lattice -family ecp5 -abc9 -run :map_ram
            opt -full
            clean -purge
            write_verilog -sv {yosys_preprocessed}
            write_rtlil {yosys_preprocessed_rtl}
            show -format dot -colors 1 -stretch -viewer none -prefix {yosys_preprocessed_graph}

            design -pop
            synth_lattice -family ecp5 -abc9 -no-rw-check
            clean -purge
            write_json {yosys_json}
            write_verilog -sv {yosys_compiled}
            tee -o {yosys_report} stat
        """))
    c.run(f"yosys -q -L {yosys_log} -T {yosys_script}")
    with open(yosys_report) as f:
        ls = f.readlines()[3:]
        print("".join(ls))
    print(f"    JSON : {yosys_json}")
    print(f"     Log : {yosys_log}")
    print(f"    Flat : {yosys_preprocessed}")
    print(f"Compiled : {yosys_compiled}")
    print(f"     Dot : {yosys_preprocessed_graph}.dot")

    pin_map = target.parent / "pin_map.lpf"
    if not pin_map.is_file():
        print(f"WARNING: no pin map at {pin_map}, skipping place&route and bitstream generation")
        return

    phase("Place and route")
    nextpnr_out = out_nextpnr / module_name.with_suffix(".pnr")
    nextpnr_log = out_nextpnr / module_name.with_suffix(".log")
    nextpnr_timing = out_nextpnr / module_name.with_suffix(".log")
    out = c.run(f"nextpnr-ecp5 --85k --detailed-timing-report -l {nextpnr_log} --report {nextpnr_timing} --json {yosys_json} --lpf {pin_map} --package=CABGA381 --speed=6 --textcfg {nextpnr_out}", hide='stderr')

    print_filtered_paragraphs(out.stderr, "Device utilisation", "Critical path", "Max frequency", "Max delay", common_prefix="Info: ")
    print(f"   PNR : {nextpnr_out}")
    print(f"   Log : {nextpnr_log}")
    print(f"Timing : {nextpnr_timing}")

    phase("Bitstream generation")
    bitstream = nextpnr_out.with_suffix(".bit")
    c.run(f"ecppack {nextpnr_out} {bitstream}")
    print(f"Wrote bitstream to {bitstream}")

@task
def test(c, target):
    for target in expand_test_target(target):
        out_info, out_sim, out_bsc = ensure_build_dirs(target, "info", "sim", "bsc")
        c.run(f"bsc -show-schedule -aggressive-conditions -check-assert -u -sim -info-dir {out_info} -simdir {out_sim} -bdir {out_bsc} -g mkTB -p {target.parent}:lib:%/Libraries {target}")
        exec = out_sim / "TB"
        c.run(f"bsc -p {out_bsc} -sim -simdir {out_sim} -e mkTB -o {exec}")
        testdata = out_sim / "testdata"
        testdata_tgt = target.parent / "testdata"
        testdata_tgt = testdata_tgt.relative_to(out_sim, walk_up=True)
        testdata.unlink(missing_ok=True)
        testdata.symlink_to(testdata_tgt, target_is_directory=True)
        with c.cd(out_sim):
            c.run(f"./TB -V {target.stem}.vcd")

@task
def clean(c):
    if Path("out").is_dir():
        shutil.rmtree("out")