#!/usr/bin/env -S python3 -m dg
import "/nmigen/cli/main"
import "/nmigen/sim/pysim/Simulator"
import "/nmigen_dg/*"
import "/os/path" qualified
import "../common/deduce_ports"


RleXorPrevDecoder = subclass Elaboratable where
    __init__ = w h data ~> None where
        @w, @h, @data, @size = w, h, data, w*h

        @pixel_x     = Signal$ range w # pixel being output
        @pixel_y     = Signal$ range h # pixel being output
        @pixel_data  = Signal! # data of pixel being output
        @pixel_valid = Signal! # wether  pixel being output is valid

        @next        = Signal! # pulse this to start decoding next frame

    elaborate = platform ~> m where with m = Module! =>
        # Default output
        Sync$ @pixel_valid ::= 0

        # Data to decode
        data_mem = Memory # TODO: use spi flash instead
            width: 8
            depth: (len @data)
            init:  (list$ bytearray$ @data)
        rddata = Submodule.rddata$ data_mem.read_port!

        # RLE decoder state:
        pixel      = Signal 1
        pixel_prev = Signal 1
        n_times    = Signal 14

        # XOR Frame buffer
        frame  = Memory
            width: 1
            depth: @size
        rdframe = Submodule.rdframe$ frame.read_port!
        wrframe = Submodule.wrframe$ frame.write_port!
        offset  = Signal$ range @size
        Comb$ rdframe.addr ::= offset
        Comb$ wrframe.addr ::= offset
        Sync$ wrframe.en   ::= 0

        do_write_pixel = Signal!
        pixel_to_write = Signal!
        Comb$ do_write_pixel ::= 0 # TODO: needed?
        When do_write_pixel $ ->
            Sync$ @pixel_data   ::= (rdframe.data ^ pixel_to_write)
            Sync$ wrframe.data  ::= (rdframe.data ^ pixel_to_write)
            Sync$ @pixel_valid  ::= 1
            Sync$ wrframe.en    ::= 1

            Sync$ offset ::= offset + 1 # needs to start at -1

            Sync$ @pixel_y ::= @pixel_y + 1
            When (@pixel_y == @h - 1) $ ->
                Sync$ @pixel_y ::= 0
                Sync$ @pixel_x ::= @pixel_x + 1

        # RLE & XOR Decoder
        FSM
            DECODED ->
                When @next $ ->
                    Sync$ offset   ::= -1 # max int
                    Sync$ @pixel_y ::= -1 # max int
                    Sync$ @pixel_x ::= 0
                    DECODED |>. READ

            READ ->
                Sync$ pixel       ::= rddata.data !! 7
                Sync$ n_times     ::= rddata.data !! slice 0 6 # TODO: proper slicing syntax would be nice
                Sync$ rddata.addr ::= rddata.addr + 1
                When
                    rddata.data !! 6 ,-> READ |>. READ_MORE
                    otherwise        ,-> READ |>. WRITE_PIXELS_BEGIN

            READ_MORE ->
                Sync$ n_times     ::= ((n_times << 8) | rddata.data)
                Sync$ rddata.addr ::= rddata.addr + 1
                READ_MORE |>. WRITE_PIXELS_BEGIN

            WRITE_PIXELS_BEGIN ->
                is_first     = offset == 0
                has_changed  = pixel  != pixel_prev
                When (~has_changed & ~is_first) $ -> # there was an omitted n_values=1 in between
                    Comb$ pixel_to_write ::= ~pixel
                    Comb$ do_write_pixel ::= 1


                Sync$ pixel_prev ::= pixel
                WRITE_PIXELS_BEGIN |>. WRITE_PIXELS

            WRITE_PIXELS ->
                Comb$ pixel_to_write ::= pixel
                Comb$ do_write_pixel ::= 1

                Sync$ n_times ::= n_times - 1
                When (n_times == 1) $ -> When # == 0 next cycle
                    offset == @size - 1 ,->  WRITE_PIXELS |>. DECODED
                    otherwise           ,->  WRITE_PIXELS |>. READ


BadAppleDecoder = -> RleXorPrevDecoder 85 64 bad_apple_data where
    fname = os.path.join
        os.path.dirname __file__
        "data"
        "bad_apple.bin"
    with f = open fname "rb" =>
        bad_apple_data = f.read!
    bad_apple_data = bad_apple_data !! slice 0 ((len bad_apple_data) // 58) # TODO: use full movie when we've moved data to spi flash
    print ((len bad_apple_data) // 58)
    #bad_apple_data = b"0000"


if __name__ == "__main__" =>
    design = RleXorPrevDecoder 85 64 b"1234"
    ports = deduce_ports design
    main design
        name: "bad_apple"
        ports: ports