# %%
import json
from amaranth import Signal
from dataclasses import dataclass
from amaranth.lib.wiring import Signature, Interface, Out
from amaranth_soc import wishbone
from amaranth.sim import Simulator
from amaranth import Elaboratable, Module

# %%
witness = json.loads(r"""
{
  "format": "Yosys Witness Trace",
  "generator": "smtbmc",
  "clocks": [
    {"path": ["\\clock"], "edge": "posedge", "offset": 0}
  ],
  "signals": [
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3532"], "width": 1, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3534"], "width": 1, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3536"], "width": 1, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3538"], "width": 1, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3540"], "width": 1, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3542"], "width": 1, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3544"], "width": 1, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3546"], "width": 1, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3548"], "width": 1, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3550"], "width": 1, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3552"], "width": 1, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3554"], "width": 32, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3556"], "width": 6, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3558"], "width": 5, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3560"], "width": 1, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3562"], "width": 5, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3564"], "width": 5, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3566"], "width": 5, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3568"], "width": 5, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3570"], "width": 32, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3572"], "width": 1, "offset": 0, "init_only": false},
    {"path": ["\\_witness_", "\\anyseq_auto_setundef_cc_533_execute_3574"], "width": 8, "offset": 0, "init_only": false},
    {"path": ["\\clock"], "width": 1, "offset": 0, "init_only": false},
    {"path": ["\\reset"], "width": 1, "offset": 0, "init_only": false},
    {"path": ["\\bus__ack"], "width": 1, "offset": 0, "init_only": false},
    {"path": ["\\uut", "\\bus__dat_r"], "width": 32, "offset": 0, "init_only": false},
    {"path": ["\\bar"], "width": 64, "offset": 0, "init_only": true},
    {"path": ["\\foo"], "width": 5, "offset": 0, "init_only": true}
  ],
  "steps": [
    {"bits": "0100000000000000000000000000000000000000000000000000000000000000000010000000000000000000000000000000001000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"},
    {"bits": "0011000000001111111111111111111100000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"},
    {"bits": "0011000000001111111111111111111100000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"},
    {"bits": "0011000000011111111111111111111100000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"},
    {"bits": "0011000000011111111111111111111100000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"},
    {"bits": "1111010001000010011111111111111100000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"},
    {"bits": "0011000000000010001101000111001110000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"},
    {"bits": "0000000000000000001000000111001100000000000010110000000000000101000001110100000000000000000000000000000000001111111111111111111111111111111000000000000"},
    {"bits": "0011001000010001001000100111001110001010101101100110010000000000100000000000000000000000000000000000001001110011111111111111111111111111111000000000000"},
    {"bits": "0011000010000001010000000001001100001001011100000000000000000000000000010000000000000000000000000000000011000000000000000000000000000000000000000000000"},
    {"bits": "0000000010010001101101010011001100010011111100000000000000000000000000000000000000000000000000000000000110000000010000000000000000000000000000000000000"},
    {"bits": "0011010000101100000000011010001100000000000100000000000000000000000000011000000000000000000000000000000000000000000000000000000000000000000100000000000"},
    {"bits": "0011010000001001101000010111001100000111100100000000000000000000000000001000000000000000000000000000000000001000000000000000000000000000000100000000000"},
    {"bits": "0000000000000000001001000000001100000111110100000000000000000000000000000000000000000000000000000000001011000000000000000000000000000000000000000000000"},
    {"bits": "0011000000100000000000000110111100010110110000000000000000000000000000000000000000000000000000000000000110011101000000000001101000000110110000000000000"},
    {"bits": "0000000000010000111000011111001100000000000100000000000000000000000000000000000000000000000000000000000100000000000000000000000000000000000000000000000"},
    {"bits": "0100010000100000010100110110001100000000000000000000000000000000000000000000000000000000000000000000000000100000000000000000000000000000000100000000000"},
    {"bits": "0000001010000100011100001001001110000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"},
    {"bits": "1011000000100000100101111110111100000011000000000000000000000000000001101100000000000000000000000000001001111000000000000000000000000000000100000000000"},
    {"bits": "0000000010000100110100001001010010001011000100000000000000000000000000000110000000000000000000000000000000011111111111111111111110100000000100000000000"},
    {"bits": "0011000001000000100000010000111100000011000000000000000000000000000000000000000000000000000000000000000011100000000000000000000000000000000000000000000"},
    {"bits": "0000000001100100000100001011011100000111010100000000000000000000000000000100000000000000000000000000000000000000000000000000000000000111001000000000000"},
    {"bits": "1111111011011000001100001001110000000111011100000000000000000000000000000000000000000000000000000000001011000000010000000000000000000000100100000000000"},
    {"bits": "1101110101010010100011111001001100000000000100000000000000000000000000000000000000000000000000000000000001100000000000000000000000000000000000000000000"},
    {"bits": "0001000001010011100000000110110100010001110000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"},
    {"bits": "0000000000000000000000000000000010000000001100000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"}
  ]
}
""")

# %%


def find_signature_1(obj, path, count=0):
    if count > 100:
        raise RecursionError

    for i, p in enumerate(path):
        sig = obj.signature
        for m in obj.signature.members:
            if m in p:
                if sig.members[m].is_signature:
                    subobj = getattr(obj, m)
                    return find_signature_1(subobj, path[i:], count + 1)
                else:
                    return getattr(obj, m)


@dataclass
class WitnessSignal:
    sig: Signal
    width: int
    start: int
    init_only: bool


def find_signature(obj: Interface, w):
    ret = []
    start = 0
    for s in w["signals"]:
        path = s["path"]

        if path in ("\\clock", "\\reset"):
            continue

        ret_1 = find_signature_1(obj, path, count=0)
        if ret_1 is not None:
            ret.append(WitnessSignal(sig=ret_1,
                                     width=s["width"],
                                     start=start,
                                     init_only=s["init_only"]))

        start += s["width"]

    return ret



# %%

def create_tb(obj, w):
    def extract_int(b, w):
        rev = b["bits"][::-1]
        offs = rev[w.start:w.start+w.width]
        rev_rev = offs[::-1]
        return int(rev_rev)

    sigs = find_signature(obj, w)

    steps = []
    for i, w_step in enumerate(w["steps"]):
        values_at_step = []
        for w_sig in sigs:
            val = extract_int(w_step, w_sig)
            values_at_step.append(w_sig.sig.eq(val))

        steps.append(values_at_step)

    def gen():
        for s in steps:
            for v in s:
                print(v)
                yield v
            yield

    return gen


# %%


class TestElaboratable(Elaboratable):
    signature = Signature({
        "bus": Out(wishbone.Signature(addr_width=30, data_width=32,
                                        granularity=8))
    })

    def __init__(self):
        self.bus = wishbone.Interface(addr_width=30, data_width=32,
                                        granularity=8)

    def elaborate(self, plat):
        m = Module()

        dummy = Signal()
        m.d.sync += dummy.eq(dummy)
        return m


m = TestElaboratable()
sim = Simulator(m)
sim.add_clock(1e-6)
sim.add_sync_process(create_tb(m, witness))

with sim.write_vcd("foo.vcd"):
    sim.run()

# %%