from migen import *

from litex.soc.interconnect.wishbone import *

from math import ceil, log2
from typing import List

from .sampler import Sampler
from .circular_buffer import CircularBuffer
from .peak_detector import PeakDetector

class SamplerController(Module):
    """
    Sampler control

    Attributes
    ----------
    bus:
        Slave wishbone bus to be connected to a higher-level bus. Has an address width set according to
        the provided buffer length.

    buffers:
        List of FIFO buffer objects used to store sample data.

    samplers:
        List of sampler objects provided by user.

    Registers
    --------
    0x00: Control Register (RW)
        Bit 0 - Begin capture. Resets all FIFOs and starts the peak detector

    0x01: Status Register (RO)
        Bit 0 - Capture complete. Set by peak detection block and cleared by software or when 

    0x02: trigger_run_len (RW)
        Number of samples to acquire after triggering sample.

    0x03: thresh_value (RW)
        Minimum peak to peak value considered triggered

    0x04: thresh_time (RW)
        Number of consecutive samples above threshold required to consider triggered

    0x05: decay_value (RW)
        Decay value to subtract from peak values to prevent false triggers

    0x06: decay_period (RW)
        Number of samples between each application of decay

    0x1xx: BUFFER_LEN_X (RO)
        Lenght of data in buffer, up to the number of samplers provided.

    """
    def __init__(self, samplers: List[Sampler], buffer_len):

        self.samplers = samplers
        num_channels = len(samplers)

        # Enables reading in samples
        sample_enable = Signal()
        # Pull in only one CDC sync signal
        sample_ready = self.samplers[0].valid

        # Generate buffers for each sampler
        self.buffers = [CircularBuffer(9, buffer_len) for _ in range(num_channels)]

        # Connect each buffer to each sampler
        for buffer, sampler in zip(self.buffers, self.samplers):
            self.comb += [
                # Connect only top 9 bits to memory
                buffer.wr_data.eq(sampler.data[1:]),
                # Writes enter FIFO only when enabled and every clock cycle
                buffer.wr_valid.eq(sample_enable & sample_ready),
            ]


        # Each sampler gets some chunk of memory at least large enough to fit
        # all of it's data, so use that as a consistent offset
        sample_mem_addr_width = ceil(log2(buffer_len))
        # 1 control block + number of channels used = control bits
        control_block_addr_width = ceil(log2(num_channels + 1))

        # Bus address width
        addr_width = control_block_addr_width + sample_mem_addr_width

        # "Master" bus
        self.bus = Interface(data_width=32, addr_width=addr_width)

        # Wishbone bus used for mapping control registers
        self.control_regs_bus = Interface(data_width=32, addr_width=sample_mem_addr_width)

        slaves = []
        slaves.append((lambda adr: adr[sample_mem_addr_width:] == 0, self.control_regs_bus))

        for i, buffer in enumerate(self.buffers):
            # Connect subordinate buses of buffers to decoder
            slaves.append((lambda adr: adr[sample_mem_addr_width:] == i + 1, buffer.bus))

            adr = (i + 1) << sample_mem_addr_width
            print(f"Sampler {i} available at 0x{adr:08x}")

        self.decoder = Decoder(self.bus, slaves)
        # TODO how to submodule
        self.submodules.decoder = self.decoder

        self.peak_detector = PeakDetector(10)
        self.comb += [
            # Simply enable whenever we start capturing
            self.peak_detector.enable.eq(sample_enable),
            # Connect to the first ADC
            self.peak_detector.data.eq(self.samplers[0].data),
            # Use the same criteria as the fifo buffer
            self.peak_detector.data_valid.eq(sample_enable & sample_ready),
        ]

        #### Control register logic

        # Storage
        control_register = Signal(32)
        status_register = Signal(32)
        trigger_run_len = Signal(32)

        def rw_register(storage: Signal, *, read: bool = True, write: bool = True):
            if read:
                read = self.control_regs_bus.dat_r.eq(storage)
            else:
                read = self.control_regs_bus.ack.eq(0)

            if write:
                write = storage.eq(self.control_regs_bus.dat_w)
            else:
                write = self.control_regs_bus.ack.eq(0)

            return If(self.control_regs_bus.we, write).Else(read)

        # Handle explicit config registers
        cases = {
            0: rw_register(control_register),
            1: rw_register(status_register, write=False),
            2: rw_register(trigger_run_len),
            3: rw_register(self.peak_detector.thresh_value),
            4: rw_register(self.peak_detector.thresh_time),
            5: rw_register(self.peak_detector.decay_value),
            6: rw_register(self.peak_detector.decay_period),

            "default": rw_register(None, read=False, write=False)
        }

        # Handle length values for each sample buffer
        for i, buffer in enumerate(self.buffers):
            cases.update({0x100 + i: rw_register(buffer.len, write=False)})

        # Connect up control registers bus
        self.sync += [
            self.control_regs_bus.ack.eq(0),
            If(self.control_regs_bus.cyc & self.control_regs_bus.stb,
                        self.control_regs_bus.ack.eq(1),
                        Case(self.control_regs_bus.adr, cases)),
        ]

        # Handle the control logic
        post_trigger_count = Signal(32)
        self.sync += [
            # Reset state whenever sampling is disabled
            If(~sample_enable, post_trigger_count.eq(0)),

            # Reset triggering status if we have started sampling
            # (peak_detector.triggered resets if sample_enable is de-asserted, so
            # this is a reliable reset mechanism)
            If(sample_enable & ~self.peak_detector.triggered,
               status_register[0].eq(0)),

            # Keep sampling past the trigger for the configured number of samples
            If(self.peak_detector.triggered & sample_enable & sample_ready,
               post_trigger_count.eq(post_trigger_count + 1),

               # We have sampled enough, update status and stop sampling
               If(post_trigger_count + 1 >= trigger_run_len,
                  status_register[0].eq(1),
                  control_register[0].eq(0))),
        ]

        # Update register storage
        self.comb += [
            sample_enable.eq(control_register[0]),
        ]

def write_wishbone(bus, address, value):
    # Set up bus
    (yield bus.adr.eq(address))
    (yield bus.dat_w.eq(value))
    (yield bus.stb.eq(1))
    (yield bus.cyc.eq(1))
    (yield bus.we.eq(1))
    yield

    cycles = 0
    while True:
        cycles += 1
        assert cycles < 5, "Write fail"


        if (yield bus.ack) == 1:
            # We received a response, clear out bus status and exit
            (yield bus.stb.eq(0))
            (yield bus.cyc.eq(0))
            yield

            break
        else:
            # Tick until we receive an ACK
            yield


def read_wishbone(bus, address,):
    """Sets up a read transaction. Due to limitations of the simulation method, you have to read
    from dat_r, and also tick immediately after calling"""
    # Set up bus
    (yield bus.adr.eq(address))
    (yield bus.stb.eq(1))
    (yield bus.cyc.eq(1))
    (yield bus.we.eq(0))
    yield

    cycles = 0
    while True:
        cycles += 1
        assert cycles < 5, "Write fail"
        if (yield bus.ack) == 1:
            # We received a response, clear out bus status and exit
            (yield bus.stb.eq(0))
            (yield bus.cyc.eq(0))

            break
        else:
            # Tick until we receive an ACK
            yield

class MockSampler(Module):
    """
    Attributes
    ----------
    All Sampler attributes by default, plus the following:

    index:
        Index of data to use from provided data
    """
    def __init__(self, data: List[int]):
        memory = Memory(width=10, depth=len(data), init=data)

        self.index = Signal(ceil(log2(len(data))))
        self.data = Signal(10)
        self.valid = Signal()

        read_port = memory.get_port(async_read=True)
        self.comb += [
            read_port.adr.eq(self.index),
            self.data.eq(read_port.dat_r),
        ]

class TestSoC(Module):
    def __init__(self, data):
        sampler = MockSampler(data)
        self.submodules.sampler = sampler
        # TODO multiple mock samplers to test that functionality
        self.controller = SamplerController([MockSampler(data)], 1024)
        self.submodules.controller = self.controller
        self.bus = self.controller.bus


def test_bus_access():
    dut = TestSoC([2, 3, 4, 5])
    def test_fn():
        yield from write_wishbone(dut.bus, 2, 0xDEADBEEF)
        yield from read_wishbone(dut.bus, 2)
        assert (yield dut.bus.dat_r) == 0xDEADBEEF, "Read failed!"

    # TODO test writing to RO register fails

    run_simulation(dut, test_fn(), vcd_name="test_bus_access.vcd")