#!/usr/bin/env python3

import re
import sys
import pprint

pp = pprint.PrettyPrinter(indent=4)

prefix = sys.argv[1]

state_types = dict()
blocks = list()
ids = dict()

def rewrite_cpp(s):
    t = list()
    i = 0
    while i < len(s):
        if s[i] in ("'", '"') and i + 1 < len(s):
            j = i + 1
            while j + 1 < len(s) and s[j] != s[i]:
                if s[j] == '\\' and j + 1 < len(s):
                    j += 1
                j += 1
            t.append(s[i:j+1])
            i = j + 1
            continue

        if s[i] in ('$', '\\') and i + 1 < len(s):
            j = i + 1
            while True:
                if j == len(s):
                    j -= 1
                    break
                if ord('a') <= ord(s[j]) <= ord('z'):
                    j += 1
                    continue
                if ord('A') <= ord(s[j]) <= ord('Z'):
                    j += 1
                    continue
                if ord('0') <= ord(s[j]) <= ord('9'):
                    j += 1
                    continue
                if s[j] == '_':
                    j += 1
                    continue
                j -= 1
                break

            n = s[i:j+1]
            i = j + 1

            if n[0] == '$':
                v = "id_d_" + n[1:]
            else:
                v = "id_b_" + n[1:]

            if v not in ids:
                ids[v] = n
            else:
                assert ids[v] == n

            t.append(v)
            continue

        if s[i] == "\t":
            t.append("  ")
        else:
            t.append(s[i])

        i += 1

    return "".join(t)

with open("%s.pmg" % prefix, "r") as f:
    while True:
        line = f.readline()
        if line == "": break
        line = line.strip()

        cmd = line.split()
        if len(cmd) == 0: continue
        cmd = cmd[0]
        
        if cmd == "state":
            m = re.match(r"^state\s+<(.*?)>\s+(([A-Za-z_][A-Za-z_0-9]*\s+)*[A-Za-z_][A-Za-z_0-9]*)\s*$", line)
            assert m
            type_str = m.group(1)
            states_str = m.group(2)
            for s in re.split(r"\s+", states_str):
                assert s not in state_types
                state_types[s] = type_str
            continue

        if cmd == "match":
            block = dict()
            block["type"] = "match"

            line = line.split()
            assert len(line) == 2
            assert line[1] not in state_types
            block["cell"] = line[1]
            state_types[line[1]] = "Cell*";

            block["select"] = list()
            block["filter"] = list()
            block["optional"] = False

            while True:
                l = f.readline()
                assert l != ""
                a = l.split()
                if len(a) == 0: continue
                if a[0] == "endmatch": break

                if a[0] == "select":
                    b = l.lstrip()[6:]
                    block["select"].append(rewrite_cpp(b.strip()))
                    continue

                if a[0] == "filter":
                    m = re.match(r"^\s*filter\s+<(.*?)>\s+(.*?)\s*===\s*(.*?)\s*$", l)
                    assert m
                    block["filter"].append((m.group(1), rewrite_cpp(m.group(2)), rewrite_cpp(m.group(3))))
                    continue

                if a[0] == "optional":
                    block["optional"] = True
                    continue

                assert False

            blocks.append(block)

        if cmd == "code":
            block = dict()
            block["type"] = "code"
            block["code"] = list()
            block["states"] = set()

            for s in line.split()[1:]:
                assert s in state_types
                block["states"].add(s)

            while True:
                l = f.readline()
                assert l != ""
                a = l.split()
                if len(a) == 0: continue
                if a[0] == "endcode": break

                block["code"].append(rewrite_cpp(l.rstrip()))

            blocks.append(block)

# pp.pprint(blocks)

with open("%s_pm.h" % prefix, "w") as f:
    print("// Generated by pmgen.py from {}.pgm".format(prefix), file=f)
    print("#include \"kernel/yosys.h\"", file=f)
    print("#include \"kernel/sigtools.h\"", file=f)
    print("YOSYS_NAMESPACE_BEGIN", file=f)
    print("struct {}_pm {{".format(prefix), file=f)
    print("  Module *module;", file=f)
    print("  SigMap sigmap;", file=f)
    print("  std::function<void(struct {}_pm*)> on_accept;".format(prefix), file=f)
    print("", file=f)

    for index in range(len(blocks)):
        block = blocks[index]
        if block["type"] == "match":
            index_types = list()
            for filt in block["filter"]:
                index_types.append(filt[0])
            print("  typedef std::tuple<{}> index_{}_key_type;".format(", ".join(index_types), index), file=f)
            print("  dict<index_{}_key_type, vector<Cell*>> index_{};".format(index, index), file=f)
    print("  pool<Cell*> blacklist;", file=f)
    print("", file=f)

    print("  struct state_t {", file=f)
    
    for s, t in sorted(state_types.items()):
        print("    {} {};".format(t, s), file=f)
    print("  } st;", file=f)
    print("", file=f)

    for v, n in sorted(ids.items()):
        if n[0] == "\\":
            print("  IdString {}{{\"\\{}\"}};".format(v, n), file=f)
        else:
            print("  IdString {}{{\"{}\"}};".format(v, n), file=f)
    print("", file=f)

    print("  {}_pm(Module *module) :".format(prefix), file=f)
    print("      module(module), sigmap(module) {", file=f)
    print("  }", file=f)

    print("", file=f)
    print("  void run(std::function<void(struct {}_pm*)> on_accept_f) {{".format(prefix), file=f)
    print("    on_accept = on_accept_f;", file=f)
    if len(blocks):
        print("    block_0();", file=f)
    print("  }", file=f)

    for index in range(len(blocks)):
        block = blocks[index]

        print("", file=f)
        print("  void block_{}() {{".format(index), file=f)

        const_st = set()
        nonconst_st = set()
        restore_st = set()

        for i in range(index):
            if blocks[i]["type"] == "code":
                for s in blocks[i]["states"]:
                    const_st.add(s)
            elif blocks[i]["type"] == "match":
                const_st.add(blocks[i]["cell"])
            else:
                assert False

        if block["type"] == "code":
            for s in block["states"]:
                if s in const_st:
                    const_st.remove(s)
                    restore_st.add(s)
                nonconst_st.add(s)
        elif block["type"] == "match":
            s = block["cell"]
            assert s not in const_st
            nonconst_st.add(s)
        else:
            assert False

        for s in sorted(const_st):
            t = state_types[s]
            if t.endswith("*"):
                print("    {} const &{} YS_ATTRIBUTE(unused) = st.{};".format(t, s, s), file=f)
            else:
                print("    const {} &{} YS_ATTRIBUTE(unused) = st.{};".format(t, s, s), file=f)

        for s in sorted(nonconst_st):
            t = state_types[s]
            print("    {} &{} YS_ATTRIBUTE(unused) = st.{};".format(t, s, s), file=f)

        if len(restore_st):
            print("", file=f)
            for s in sorted(restore_st):
                t = state_types[s]
                print("    {} backup_{} = st.{};".format(t, s, s), file=f)

        if len(restore_st):
            print("", file=f)
            for s in sorted(restore_st):
                t = state_types[s]
                print("    st.{} = backup_{};".format(s, s), file=f)

        print("  }", file=f)

    print("};\nYOSYS_NAMESPACE_END", file=f)