from sim_client import SimBusClient

from amaranth import Module, Signal

from amaranth.sim import Simulator, SimulatorContext

from amaranth_ila.ila import ILACapture, ILAControl
from amaranth_ila.client import csr

def test_ila_capture():
    m = Module()

    cnt = Signal(8)
    m.d.sync += cnt.eq(cnt + 1)

    m.submodules.capture = capture = ILACapture(width = 8, depth = 64)
    m.submodules.control = control = ILAControl(capture)

    m.d.comb += [
        capture.inputs.eq(cnt),
        capture.trigger.eq(cnt == 18),
        capture.pre_trigger.eq(4),
    ]

    sim = Simulator(m)
    sim.add_clock(1e-6)

    @sim.add_testbench
    async def testbench(ctx: SimulatorContext):
        client = csr.Client(control.metadata.as_json()['interface']['members']['bus']['annotations'], bus_client = SimBusClient(ctx, control.bus))

        # Start the capture.
        await client.cmd.write(1)

        # Wait for the capture to finish.
        while await client.status.read() != 3:
            pass

        # Read the captured data.
        data = [await client.data.read() for _ in range(64)]

        # Confirm all data is drained.
        await ctx.tick()
        assert ctx.get(capture.output.valid) == 0

        assert data == list(range(18 - 4, 18 - 4 + 64))

    @sim.add_process
    async def timeout(ctx: SimulatorContext):
        await ctx.tick().repeat(10_000)
        raise TimeoutError('Simulation timed out')

    with sim.write_vcd('test.vcd'):
        sim.run()