use core::mem::MaybeUninit;

use alloc::sync::Arc;
use device_api::{
    device::{Device, DeviceInitContext},
    dma::DmaAllocator,
    interrupt::{InterruptHandler, IrqVector},
};
use libk::{dma::DmaBuffer, error::Error};
use libk_mm::{address::PhysicalAddress, device::DeviceMemoryIo};
use libk_util::{queue::BoundedQueue, sync::IrqSafeSpinlock, OneTimeInit};
use tock_registers::{
    interfaces::{ReadWriteable, Readable, Writeable},
    register_bitfields, register_structs,
    registers::{ReadOnly, ReadWrite},
};
use ygg_driver_net_core::{
    interface::{NetworkDevice, NetworkInterfaceType},
    RxPacket,
};
use ygg_driver_pci::device::PciDeviceInfo;
use yggdrasil_abi::net::MacAddress;

register_bitfields! {
    u8,

    pub CR [
        /// Reset. Setting this bit forces the RTL8139D(L) to perform a software reset. This bit is
        /// cleared by the RTL8139D(L) after the reset completes.
        RST     OFFSET(4)   NUMBITS(1)  [],
        /// Receiver enable
        RE      OFFSET(3)   NUMBITS(1)  [],
        /// Transmitter enable
        TE      OFFSET(2)   NUMBITS(1)  [],
        /// Rx buffer empty
        BUFE    OFFSET(0)   NUMBITS(1)  [],
    ],

    pub CONFIG1 [
        /// LED pin control
        LEDS1   OFFSET(7)   NUMBITS(1)  [],
        /// LED pin control
        LEDS0   OFFSET(6)   NUMBITS(1)  [],
        /// Driver load. Software may use this bit to make sure that the driver has been loaded.
        DVRLOAD OFFSET(5)   NUMBITS(1)  [],
        /// LWAKE active mode. The LWACT bit and LWPTN bit in CONFIG4 are used to program the LWAKE
        /// pin's output signal.
        LWACT   OFFSET(4)   NUMBITS(1)  [],
        /// The operational registers are mapped into PCI memory space
        MEMMAP  OFFSET(3)   NUMBITS(1)  [],
        /// The operational registers are mapped into PCI I/O space
        IOMAP   OFFSET(2)   NUMBITS(1)  [],
        /// Set to enable Vital Product Data
        VPD     OFFSET(1)   NUMBITS(1)  [],
        /// Power management enable
        PMEn    OFFSET(0)   NUMBITS(1)  [],
    ],
}

register_bitfields! {
    u16,

    pub IMR_ISR [
        /// System error interrupt
        SERR        OFFSET(15)  NUMBITS(1)  [],
        /// Time out interrupt
        TIMEOUT     OFFSET(14)  NUMBITS(1)  [],
        /// Cable length change interrupt
        LENCHG      OFFSET(13)  NUMBITS(1)  [],
        /// Rx FIFO overflow interrupt
        FOVW        OFFSET(6)   NUMBITS(1)  [],
        /// Packet underrun/Link change interrupt
        PUN_LINKCHG OFFSET(5)   NUMBITS(1)  [],
        /// Rx buffer overflow
        RXOVW       OFFSET(4)   NUMBITS(1)  [],
        /// Tx error interrupt
        TER         OFFSET(3)   NUMBITS(1)  [],
        /// Tx OK interrupt
        TOK         OFFSET(2)   NUMBITS(1)  [],
        /// Rx error interrupt
        RER         OFFSET(1)   NUMBITS(1)  [],
        /// Rx OK interrupt
        ROK         OFFSET(0)   NUMBITS(1)  [],
    ],
}

register_bitfields! {
    u32,

    pub TSDn [
        /// Carrier sense lost
        CRS         OFFSET(31)  NUMBITS(1)  [],
        /// Tx aborted
        TABT        OFFSET(30)  NUMBITS(1)  [],
        /// Out of window collision
        OWC         OFFSET(29)  NUMBITS(1)  [],
        /// CD heart beat
        CDH         OFFSET(28)  NUMBITS(1)  [],
        /// Number of collision count
        NCC         OFFSET(24)  NUMBITS(4)  [],
        /// Early Tx threshold
        ERTXTH      OFFSET(16)  NUMBITS(6)  [],
        /// Tx OK
        TOK         OFFSET(15)  NUMBITS(1)  [],
        /// Tx FIFO underrun
        TUN         OFFSET(14)  NUMBITS(1)  [],
        /// Tx descriptor is owned by the DMA. Set to 0 by the driver to initiate a send
        OWN         OFFSET(13)  NUMBITS(1)  [],
        /// Tx size
        SIZE        OFFSET(0)   NUMBITS(13) [],
    ],

    pub RCR [
        /// Early Rx threshold bits
        ERTH        OFFSET(24)  NUMBITS(4)  [],
        /// Multiple early interrupt select
        MULERINT    OFFSET(17)  NUMBITS(1)  [],
        /// TBD, didn't really understand the docs, lmao
        RER8        OFFSET(16)  NUMBITS(1)  [],
        /// Rx FIFO threshold
        RXFTH       OFFSET(13)  NUMBITS(3)  [],
        /// Rx buffer length (+16 bytes)
        RBLEN       OFFSET(11)  NUMBITS(2)  [
            Len8K = 0,
            Len16K = 1,
            Len32K = 2,
            Len64K = 3,
        ],
        /// Max DMA burst size per Rx DMA burst
        MXDMA       OFFSET(8)   NUMBITS(3)  [],
        /// When 0: the packet will get split if crossing the Rx buffer end
        /// When 1: the packet will get placed contiguously
        WRAP        OFFSET(7)   NUMBITS(1)  [],
        /// Accept error packet
        AER         OFFSET(5)   NUMBITS(1)  [],
        /// Accept runt (packets smaller than 64 bytes)
        AR          OFFSET(4)   NUMBITS(1)  [],
        /// Accept broadcast packets
        AB          OFFSET(3)   NUMBITS(1)  [],
        /// Accept multicast packets
        AM          OFFSET(2)   NUMBITS(1)  [],
        /// Accept physical match
        APM         OFFSET(1)   NUMBITS(1)  [],
        /// Accept all packets
        AAP         OFFSET(0)   NUMBITS(1)  [],
    ],
}

register_structs! {
    #[allow(non_snake_case)]
    pub Regs {
        (0x0000 => pub IDRn: [ReadWrite<u32>; 2]),
        (0x0008 => pub MARn: [ReadWrite<u8>; 8]),
        (0x0010 => pub TSDn: [ReadWrite<u32, TSDn::Register>; 4]),
        (0x0020 => pub TSADn: [ReadWrite<u32>; 4]),
        (0x0030 => pub RBSTART: ReadWrite<u32>),
        (0x0034 => pub ERBCR: ReadOnly<u16>),
        (0x0036 => pub ERSR: ReadOnly<u8>),
        (0x0037 => pub CR: ReadWrite<u8, CR::Register>),
        (0x0038 => pub CAPR: ReadWrite<u16>),
        (0x003A => pub CBR: ReadOnly<u16>),
        (0x003C => pub IMR: ReadWrite<u16, IMR_ISR::Register>),
        (0x003E => pub ISR: ReadWrite<u16, IMR_ISR::Register>),
        (0x0040 => pub TCR: ReadWrite<u32>),
        (0x0044 => pub RCR: ReadWrite<u32, RCR::Register>),
        (0x0048 => pub TCTR: ReadWrite<u32>),
        (0x004C => pub MPC: ReadWrite<u32>),
        (0x0050 => pub r9346CR: ReadWrite<u8>),
        (0x0051 => pub CONFIG0: ReadOnly<u8>),
        (0x0052 => pub CONFIG1: ReadWrite<u8, CONFIG1::Register>),
        (0x0053 => _1),
        (0x0054 => pub TIMERINT: ReadWrite<u32>),
        (0x0058 => pub MSR: ReadWrite<u8>),
        (0x0059 => pub CONFIG3: ReadWrite<u8>),
        (0x005A => pub CONFIG4: ReadWrite<u8>),
        (0x005B => _2),
        (0x005C => pub MULINT: ReadWrite<u16>),
        (0x005E => pub RERID: ReadOnly<u8>),
        (0x005F => _3),
        (0x0060 => pub TSAD: ReadOnly<u16>),
        (0x0062 => pub BMCR: ReadWrite<u16>),
        (0x0064 => pub BMSR: ReadOnly<u16>),
        (0x0066 => pub ANAR: ReadWrite<u16>),
        (0x0068 => pub ANLPAR: ReadOnly<u16>),
        (0x006A => pub ANER: ReadOnly<u16>),
        (0x006C => pub DIS: ReadOnly<u16>),
        (0x006E => pub FCSC: ReadOnly<u16>),
        (0x0070 => pub NWAYTR: ReadWrite<u16>),
        (0x0072 => pub REC: ReadOnly<u16>),
        (0x0074 => pub CSCR: ReadWrite<u16>),
        (0x0076 => _4),
        (0x0078 => pub PHY1_PARM: ReadWrite<u32>),
        (0x007C => pub TW_PARM: ReadWrite<u32>),
        (0x0080 => pub PHY2_PARM: ReadWrite<u8>),
        (0x0081 => _5),
        (0x0084 => pub CRCn: [ReadWrite<u8>; 8]),
        (0x008C => pub WAKEUPn: [ReadWrite<u32>; 16]),
        (0x00CC => pub LSBCRCn: [ReadWrite<u8>; 8]),
        (0x00D4 => _6),
        (0x00D8 => pub CONFIG5: ReadWrite<u8>),
        (0x00D9 => _7),
        (0x0100 => @END),
    }
}

struct Rx {
    buffer: DmaBuffer<[MaybeUninit<u8>]>,
    rd: usize,
}

// TODO place a secondary Tx queue here, to send the queued packets when more slots become
// available
struct Tx {
    buffers: [Option<DmaBuffer<[u8]>>; 4],
    wr: usize,
    rd: usize,
    queue: BoundedQueue<DmaBuffer<[u8]>>,
}

pub struct Rtl8139 {
    regs: IrqSafeSpinlock<DeviceMemoryIo<'static, Regs>>,
    mac: MacAddress,
    pci: PciDeviceInfo,

    nic: OneTimeInit<u32>,
    rx: OneTimeInit<IrqSafeSpinlock<Rx>>,
    tx: IrqSafeSpinlock<Tx>,

    dma: Arc<dyn DmaAllocator>,
}

impl Tx {
    pub fn tx_now(&mut self, regs: &Regs, packet: DmaBuffer<[u8]>) -> Result<(), Error> {
        let packet_address = packet.bus_address().try_into_u32()?;
        let packet_len = packet.len();
        if packet_len > 1500 {
            return Err(Error::InvalidArgument);
        }

        self.buffers[self.wr] = Some(packet);
        regs.TSADn[self.wr].set(packet_address);
        regs.TSDn[self.wr].write(TSDn::SIZE.val(packet_len as u32));

        self.wr = (self.wr + 1) % 4;

        Ok(())
    }

    pub fn is_full(&self) -> bool {
        (self.wr + 1) % 4 == self.rd
    }
}

impl Rtl8139 {
    // 8K + overflow space
    const RX_BUFFER_LEN: usize = 8192;
    const RX_BUFFER_OVERFLOW: usize = 4096;

    pub fn new(
        dma: Arc<dyn DmaAllocator>,
        base: PhysicalAddress,
        info: PciDeviceInfo,
    ) -> Result<Self, Error> {
        let regs = unsafe { DeviceMemoryIo::<Regs>::map(base, Default::default()) }?;
        let mac0 = regs.IDRn[0].get().to_le_bytes();
        let mac1 = regs.IDRn[1].get().to_le_bytes();
        let mac = MacAddress::from([mac0[0], mac0[1], mac0[2], mac0[3], mac1[0], mac1[1]]);

        Ok(Self {
            mac,
            regs: IrqSafeSpinlock::new(regs),
            pci: info,

            nic: OneTimeInit::new(),
            rx: OneTimeInit::new(),
            tx: IrqSafeSpinlock::new(Tx {
                buffers: [const { None }; 4],
                wr: 0,
                rd: 0,
                queue: BoundedQueue::new(64),
            }),

            dma,
        })
    }
}

impl InterruptHandler for Rtl8139 {
    fn handle_irq(self: Arc<Self>, _vector: IrqVector) -> bool {
        let regs = self.regs.lock();
        let status = regs.ISR.extract();
        // Clear ISR bits
        regs.ISR.write(IMR_ISR::TOK::SET + IMR_ISR::ROK::SET);

        let mut any = false;
        if status.matches_all(IMR_ISR::TOK::SET) {
            let mut tx = self.tx.lock();

            tx.rd = (tx.rd + 1) % 4;
            if let Some(packet) = tx.queue.pop() {
                // Should not fail, we just freed a single descriptor
                tx.tx_now(&regs, packet).unwrap();
            }

            any = true;
        }
        if status.matches_all(IMR_ISR::ROK::SET) {
            let nic = *self.nic.get();
            let mut rx = self.rx.get().lock();
            let cbr = regs.CBR.get();

            loop {
                let rx_pos = rx.rd % Self::RX_BUFFER_LEN;
                if rx_pos == cbr as usize {
                    break;
                }

                // 4-byte preamble
                let rx_len_0 = unsafe { rx.buffer[rx_pos + 2].assume_init() };
                let rx_len_1 = unsafe { rx.buffer[rx_pos + 3].assume_init() };
                let rx_len = u16::from_le_bytes([rx_len_0, rx_len_1]) as usize;

                if rx_len >= 16
                    && let Ok(mut packet_buf) = DmaBuffer::new_uninit_slice(&*self.dma, rx_len)
                {
                    packet_buf.copy_from_slice(&rx.buffer[rx_pos + 4..rx_pos + rx_len + 4]);
                    let packet_buf = unsafe { DmaBuffer::assume_init_slice(packet_buf) };
                    // let packet_buf = unsafe { packet_buf.assume_init_slice() };
                    let packet = RxPacket::new(packet_buf, 0, nic);
                    ygg_driver_net_core::receive_packet(packet).ok();
                }

                // rx_len + 4, aligned to 4 bytes
                let total_len = (rx_len + 7) & !3;
                rx.rd = rx.rd.wrapping_add(total_len);
            }

            regs.CAPR.set((rx.rd as u16).max(cbr).wrapping_sub(0x10));
            rx.rd %= Self::RX_BUFFER_LEN;

            any = true;
        }

        any
    }
}

impl Device for Rtl8139 {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        log::info!("Initialize rtl8139 driver");
        log::info!("MAC: {}", self.mac);

        // Setup the initial Rx buffer
        let rx_buffer = DmaBuffer::new_uninit_slice(
            &*self.dma,
            Self::RX_BUFFER_LEN + Self::RX_BUFFER_OVERFLOW,
        )?;
        let rx_buffer_address = rx_buffer.bus_address().try_into_u32()?;

        self.pci.map_interrupt(Default::default(), self.clone())?;

        let regs = self.regs.lock();

        // Power up the device
        regs.CONFIG1.set(0);

        // Reset the device
        let mut timeout = 1000000;
        regs.CR.write(CR::RST::SET);
        while timeout > 0 && regs.CR.matches_all(CR::RST::SET) {
            core::hint::spin_loop();
            timeout -= 1;
        }
        if timeout == 0 {
            log::error!("rtl8139: reset timed out");
            return Err(Error::TimedOut);
        }

        // Configure Rx
        regs.RBSTART.set(rx_buffer_address);
        regs.RCR
            .write(RCR::WRAP::SET + RCR::AB::SET + RCR::APM::SET + RCR::AR::SET + RCR::AM::SET);

        // Enable Rx/Tx interrupts
        regs.IMR.write(IMR_ISR::ROK::SET + IMR_ISR::TOK::SET);

        // Start Rx/Tx
        regs.CR.modify(CR::TE::SET + CR::RE::SET);

        self.rx.init(IrqSafeSpinlock::new(Rx {
            buffer: rx_buffer,
            rd: 0,
        }));
        let nic =
            ygg_driver_net_core::register_interface(NetworkInterfaceType::Ethernet, self.clone());
        self.nic.init(nic.id());

        Ok(())
    }

    fn display_name(&self) -> &str {
        "Realtek RTL8139 10/100Mbps Ethernet Controller"
    }
}

impl NetworkDevice for Rtl8139 {
    fn allocate_transmit_buffer(&self, len: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error> {
        DmaBuffer::new_uninit_slice(&*self.dma, len)
    }

    fn transmit_buffer(&self, packet: DmaBuffer<[u8]>) -> Result<(), Error> {
        let mut tx = self.tx.lock();

        // Buffer still in Tx, cannot send
        if tx.is_full() {
            if tx.queue.push(packet).is_err() {
                log::warn!("rtl8139: out of tx buffers + queue full");
                return Err(Error::WouldBlock);
            } else {
                return Ok(());
            }
        }
        let regs = self.regs.lock();
        tx.tx_now(&regs, packet)
    }

    fn packet_prefix_size(&self) -> usize {
        0
    }

    fn read_hardware_address(&self) -> MacAddress {
        self.mac
    }
}
