import logging
import asyncio

from amaranth import *
from amaranth.lib import enum, data, wiring, stream, io
from amaranth.lib.wiring import In, Out

from glasgow.abstract import AbstractAssembly, GlasgowPin
from glasgow.applet import GlasgowAppletV2

class MUComponent(wiring.Component):
    o_stream: Out(stream.Signature(8))

    def __init__(self, ports, flush):
        self._ports = ports
        self._flush = flush
        super().__init__()

    def elaborate(self, platform):
        m = Module()
        m.submodules.clk_buffer  = clkb  = io.Buffer("i", self._ports.clk)
        m.submodules.lr_buffer   = lrb   = io.Buffer("i", self._ports.lr)
        m.submodules.data_buffer = datab = io.Buffer("i", self._ports.data)

        leds = [platform.request("led", n) for n in range(5)]

        clk0 = Signal()
        clk1 = Signal()
        clk2 = Signal()
        clkp = Signal()
        lr0 = Signal()
        lr1 = Signal()
        lr2 = Signal()
        lrp = Signal()
        data0 = Signal()
        data1 = Signal()
        data2 = Signal()

        sr = Signal(128)
        srb = Signal(128+8)
        src = Signal(5)
        sc = Signal(8)

        tick = Signal()
        
        m.d.sync += [clk0.eq( clkb.i),  clk1.eq( clk0),  clk2.eq (clk1)]
        m.d.sync += [lr0.eq(  lrb.i),   lr1.eq(  lr0),   lr2.eq  (lr1)]
        m.d.sync += [data0.eq(datab.i), data1.eq(data0), data2.eq(data1)]

        m.d.comb += self.o_stream.payload.eq(srb[128:])

        m.d.comb += clkp.eq(~clk2 & clk1)
        m.d.comb += self.o_stream.valid.eq(src != 0)

        with m.If(clkp):
            m.d.sync += tick.eq(1)
            with m.If(tick):
                m.d.sync += lrp.eq(lr2)
                m.d.sync += sr.eq(Cat(data2, sr[:127]))
            
                with m.If(lr2 & ~lrp):
                    m.d.sync += [srb.eq(Cat(C(0x00112233445566778899aabbccddeeff, 128), sc)), src.eq(17)]
                    m.d.sync += sc.eq(sc+1)

        with m.If(self.o_stream.ready & self.o_stream.valid):
            m.d.sync += srb.eq(Cat(C(0, 8), srb[:128]))
            m.d.sync += src.eq(src-1)
            with m.If(src == 1):
                m.d.comb += self._flush.eq(1)

#        m.d.comb += [leds[i].o.eq(count[i+18]) for i in range(5)]

        return m

class MUInterface:
    def __init__(self, logger: logging.Logger, assembly: AbstractAssembly):
        self._logger = logger
        self._level  = logging.DEBUG if self._logger.name == __name__ else logging.TRACE
        ports = assembly.add_port_group(clk="A5", lr="A4", data="A3", datam="A6")
        flush = Signal(1)
        component = assembly.add_submodule(MUComponent(ports, flush))
        self._pipe = assembly.add_in_pipe(component.o_stream, in_flush=flush)

    def _log(self, message: str, *args):
        self._logger.log(self._level, "mu: " + message, *args)

    async def run(self):
        zz = 0
        while True:
            v = await self._pipe.recv(17)
            print("%02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x" % (v[0], v[1], v[2], v[3], v[4], v[5], v[6], v[7], v[8], v[9], v[10], v[11], v[12], v[13], v[14], v[15], v[16]))


class MUApplet(GlasgowAppletV2):
    logger = logging.getLogger(__name__)
    help = "mu capture applet"
    preview = False
    description = """
    MU capture applet.
    """

    @classmethod
    def add_setup_arguments(cls, parser):
        pass

    @classmethod
    def add_build_arguments(cls, parser, access):
        pass
    
    def build(self, args):
        with self.assembly.add_applet(self):
            self.assembly.use_voltage({'A':5.0, 'B':5.0})
            self.mu_iface = MUInterface(self.logger, self.assembly)

    @classmethod
    def add_run_arguments(cls, parser):
        pass

    async def run(self, args):
        await self.mu_iface.run()