# SPDX-License-Identifier: CERN-OHL-S-2.0
"""

Copyright (c) 2025 FPGA Ninja, LLC

Authors:
- Alex Forencich

"""

import logging
import struct
from collections import deque

from cocotb.queue import Queue

class Port:
    def __init__(self, driver, index, hw_regs):
        self.driver = driver
        self.log = driver.log
        self.index = index
        self.hw_regs = hw_regs

        self.rxq_log_size = (256).bit_length()-1
        self.rxq_size = 2**self.rxq_log_size
        self.rxq_mask = self.rxq_size-1
        self.rxq = None
        self.rxq_prod = 0
        self.rxq_cons = 0

        self.rx_info = [None] * self.rxq_size

        self.rxcq_log_size = (256).bit_length()-1
        self.rxcq_size = 2**self.rxcq_log_size
        self.rxcq_mask = self.rxcq_size-1
        self.rxcq = None
        self.rxcq_prod = 0
        self.rxcq_cons = 0

        self.txq_log_size = (256).bit_length()-1
        self.txq_size = 2**self.txq_log_size
        self.txq_mask = self.txq_size-1
        self.txq = None
        self.txq_prod = 0
        self.txq_cons = 0

        self.tx_info = [None] * self.txq_size

        self.txcq_log_size = (256).bit_length()-1
        self.txcq_size = 2**self.txcq_log_size
        self.txcq_mask = self.txcq_size-1
        self.txcq = None
        self.txcq_prod = 0
        self.txcq_cons = 0

        self.rx_queue = Queue()

    async def init(self):

        self.rxq = self.driver.pool.alloc_region(self.rxq_size*16)
        addr = self.rxq.get_absolute_address(0)
        await self.hw_regs.write_dword(0x0200, 0x00000000)
        await self.hw_regs.write_dword(0x0204, 0x00000000)
        await self.hw_regs.write_dword(0x0208, addr & 0xffffffff)
        await self.hw_regs.write_dword(0x020c, addr >> 32)
        await self.hw_regs.write_dword(0x0200, 0x00000001 | (self.rxq_log_size << 16))

        self.rxcq = self.driver.pool.alloc_region(self.rxcq_size*16)
        addr = self.rxcq.get_absolute_address(0)
        await self.hw_regs.write_dword(0x0400, 0x00000000)
        await self.hw_regs.write_dword(0x0408, addr & 0xffffffff)
        await self.hw_regs.write_dword(0x040c, addr >> 32)
        await self.hw_regs.write_dword(0x0400, 0x00000001 | (self.rxcq_log_size << 16))

        self.txq = self.driver.pool.alloc_region(self.txq_size*16)
        addr = self.txq.get_absolute_address(0)
        await self.hw_regs.write_dword(0x0100, 0x00000000)
        await self.hw_regs.write_dword(0x0104, 0x00000000)
        await self.hw_regs.write_dword(0x0108, addr & 0xffffffff)
        await self.hw_regs.write_dword(0x010c, addr >> 32)
        await self.hw_regs.write_dword(0x0100, 0x00000001 | (self.txq_log_size << 16))

        self.txcq = self.driver.pool.alloc_region(self.txcq_size*16)
        addr = self.txcq.get_absolute_address(0)
        await self.hw_regs.write_dword(0x0300, 0x00000000)
        await self.hw_regs.write_dword(0x0308, addr & 0xffffffff)
        await self.hw_regs.write_dword(0x030c, addr >> 32)
        await self.hw_regs.write_dword(0x0300, 0x00000001 | (self.txcq_log_size << 16))

        # wait for writes to complete
        await self.hw_regs.read_dword(0)

        await self.refill_rx_buffers()

    async def start_xmit(self, data):
        headroom = 10
        tx_buf = self.driver.alloc_pkt()
        await tx_buf.write(headroom, data)
        index = self.txq_prod & self.txq_mask
        ptr = tx_buf.get_absolute_address(0)
        struct.pack_into('<xxxxLQ', self.txq.mem, 16*index, len(data), ptr+headroom)
        self.tx_info[index] = tx_buf
        self.txq_prod += 1
        await self.hw_regs.write_dword(0x0104, self.txq_prod & 0xffff)

    async def recv(self):
        return await self.rx_queue.get()

    async def recv_nowait(self):
        return self.rx_queue.get_nowait()

    def free_tx_desc(self, index):
        pkt = self.tx_info[index]
        self.driver.free_pkt(pkt)
        self.tx_info[index] = None

    def free_tx_buf(self):
        while self.txq_cons != self.txq_prod:
            index = self.txq_cons & self.txq_mask
            self.free_tx_desc(index)
            self.txq_cons += 1

    async def process_tx_cq(self):

        cq_cons_ptr = self.txcq_cons
        cons_ptr = self.txq_cons

        while True:
            cq_index = cq_cons_ptr & self.txcq_mask
            index = cons_ptr & self.txq_mask

            cpl_data = struct.unpack_from("<LLLL", self.txcq.mem, cq_index*16)

            self.log.info("TX CQ index %d data %s", cq_index, cpl_data)

            if bool(cpl_data[-1] & 0x80000000) == bool(cq_cons_ptr & self.txcq_size):
                self.log.info("CQ empty")
                break

            pkt = self.tx_info[index]

            self.free_tx_desc(index)

            cq_cons_ptr += 1
            cons_ptr += 1

        self.txcq_cons = cq_cons_ptr
        self.txq_cons = cons_ptr

    def free_rx_desc(self, index):
        pkt = self.rx_info[index]
        self.driver.free_pkt(pkt)
        self.rx_info[index] = None

    def free_rx_buf(self):
        while self.rxq_cons != self.rxq_prod:
            index = self.rxq_cons & self.rxq_mask
            self.free_rx_desc(index)
            self.rxq_cons += 1

    def prepare_rx_desc(self, index):
        pkt = self.driver.alloc_pkt()
        self.rx_info[index] = pkt

        length = pkt.size
        ptr = pkt.get_absolute_address(0)

        struct.pack_into('<xxxxLQ', self.rxq.mem, 16*index, length, ptr)

    async def refill_rx_buffers(self):
        missing = self.rxq_size - (self.rxq_prod - self.rxq_cons)

        if missing < 8:
            return

        for k in range(missing):
            self.prepare_rx_desc(self.rxq_prod & self.rxq_mask)
            self.rxq_prod += 1

        await self.hw_regs.write_dword(0x0204, self.rxq_prod & 0xffff)

    async def process_rx_cq(self):

        cq_cons_ptr = self.rxcq_cons
        cons_ptr = self.rxq_cons

        while True:
            cq_index = cq_cons_ptr & self.rxcq_mask
            index = cons_ptr & self.rxq_mask

            cpl_data = struct.unpack_from("<LLLL", self.rxcq.mem, cq_index*16)

            self.log.info("RX CQ index %d data %s", cq_index, cpl_data)

            if bool(cpl_data[-1] & 0x80000000) == bool(cq_cons_ptr & self.rxcq_size):
                self.log.info("CQ empty")
                break

            pkt = self.rx_info[index]
            length = cpl_data[1]

            data = pkt[:length]

            self.log.info("Packet: %s", data)

            self.rx_queue.put_nowait(data)

            self.free_rx_desc(index)

            cq_cons_ptr += 1
            cons_ptr += 1

        self.rxcq_cons = cq_cons_ptr
        self.rxq_cons = cons_ptr

        await self.refill_rx_buffers()

    async def interrupt_handler(self):
        self.log.info("Interrupt")
        await self.process_rx_cq()
        await self.process_tx_cq()


class Driver:
    def __init__(self):
        self.log = logging.getLogger("cocotb.cndm_proto")

        self.dev = None
        self.pool = None
        self.hw_regs = None

        self.ports = []

        self.free_packets = deque()
        self.allocated_packets = []

    async def init_pcie_dev(self, dev):
        self.dev = dev
        self.pool = dev.rc.mem_pool

        await dev.enable_device()
        await dev.set_master()
        await dev.alloc_irq_vectors(32, 32)

        self.hw_regs = dev.bar_window[0]

        await self.init_common()

    async def init_common(self):
        self.port_count = await self.hw_regs.read_dword(0x0100)
        self.port_offset = await self.hw_regs.read_dword(0x0104)
        self.port_stride = await self.hw_regs.read_dword(0x0108)

        self.log.info("Port count: %d", self.port_count)
        self.log.info("Port offset: 0x%x", self.port_offset)
        self.log.info("Port stride: 0x%x", self.port_stride)

        for k in range(self.port_count):
            port = Port(self, k, self.hw_regs.create_window(self.port_offset + self.port_stride*k))
            await port.init()
            self.dev.request_irq(k, port.interrupt_handler)

            self.ports.append(port)

    def alloc_pkt(self):
        if self.free_packets:
            return self.free_packets.popleft()

        pkt = self.pool.alloc_region(4096)
        self.allocated_packets.append(pkt)
        return pkt

    def free_pkt(self, pkt):
        assert pkt is not None
        assert pkt in self.allocated_packets
        self.free_packets.append(pkt)
