#![no_std]
#![recursion_limit = "256"]

use core::{mem::MaybeUninit, time::Duration};

use alloc::sync::Arc;
use device_api::{
    clock::{ClockHandle, ResetHandle},
    device::{Device, DeviceInitContext},
    dma::DmaAllocator,
    interrupt::{InterruptHandler, IrqHandle, IrqVector},
};
use device_tree::driver::{
    device_tree_driver, util::read_mac_address, InitSequence, Node, ProbeContext,
};
use futures_util::task::AtomicWaker;
use libk::{
    dma::DmaBuffer,
    error::Error,
    task::runtime::{self, psleep, pwait},
};
use libk_mm::{address::PhysicalAddress, device::DeviceMemoryIo};
use libk_util::{event::BitmapEvent, sync::IrqSafeSpinlock, OneTimeInit};
use regs::{
    dma::{DMACiCR, DMACiIER, DMACiSR, DMACiTXCR, DMAC0RXCR, DMAMR, DMASBMR},
    mac::{
        MACAiHR, MACAiLR, MACCR, MACIER, MACISR, MACPFR, MACPHYCSR, MACQ0TXFCR, MACRXFCR, MACRXQC0R,
    },
    mtl::{MTLRXQiOMR, MTLTXQiOMR, MTLOMR},
    Regs,
};
use ring::{RxDescriptor, RxRing, TxDescriptor, TxRing};
use tock_registers::interfaces::{ReadWriteable, Readable, Writeable};
use ygg_driver_net_core::{
    ephy::{PhyAccess, GBESR},
    interface::{NetworkDevice, NetworkInterfaceType},
    RxPacket,
};
use yggdrasil_abi::net::{
    link::{Duplex, EthernetLinkState, EthernetSpeed, LinkState},
    MacAddress,
};

extern crate alloc;

pub mod regs;
pub mod ring;

struct Inner {
    regs: IrqSafeSpinlock<DeviceMemoryIo<'static, Regs>>,

    tx_ring: IrqSafeSpinlock<TxRing>,
    rx_ring: IrqSafeSpinlock<RxRing>,
}

struct Stmmac {
    base: PhysicalAddress,
    dma: OneTimeInit<Arc<dyn DmaAllocator>>,
    mac: MacAddress,

    clk_stmmaceth: ClockHandle,
    clk_pclk: ClockHandle,
    clk_ptp_ref: ClockHandle,
    clk_tx: ClockHandle,
    clk_gtx: ClockHandle,

    rst_stmmaceth: ResetHandle,
    rst_ahb: ResetHandle,

    irq: IrqHandle,
    softirq_events: BitmapEvent<AtomicWaker>,

    inner: OneTimeInit<Inner>,
    iface_id: OneTimeInit<u32>,
}

impl Inner {
    fn link_state(&self) -> LinkState {
        let status = self.regs.lock().MAC.MACPHYCSR.extract();
        let link = if status.matches_all(MACPHYCSR::LNKSTS::Up) {
            let speed = match status.read_as_enum(MACPHYCSR::LNKSPEED) {
                Some(MACPHYCSR::LNKSPEED::Value::Speed2_5MHz) => EthernetSpeed::Speed10,
                Some(MACPHYCSR::LNKSPEED::Value::Speed25MHz) => EthernetSpeed::Speed100,
                Some(MACPHYCSR::LNKSPEED::Value::Speed125MHz) => EthernetSpeed::Speed1000,
                _ => EthernetSpeed::Unknown,
            };
            let duplex = if status.matches_all(MACPHYCSR::LNKMOD::FullDuplex) {
                Duplex::Full
            } else {
                Duplex::Half
            };
            EthernetLinkState::Up(speed, duplex)
        } else {
            EthernetLinkState::Down
        };

        LinkState::Ethernet(link)
    }
}

impl Stmmac {
    const SOFTIRQ_GMII_STATUS: u64 = 1 << 0;
    const SOFTIRQ_TX_STATUS: u64 = 1 << 1;
    const SOFTIRQ_RX_STATUS: u64 = 1 << 2;

    pub fn start_xmit(&self, frame: DmaBuffer<[u8]>) -> Result<(), Error> {
        let inner = self.inner.get();
        let regs = inner.regs.lock();

        let mut tx_ring = inner.tx_ring.lock();
        let index = tx_ring.push_xmit(frame)?;
        let ring_pos = tx_ring
            .buffer_base()
            .add(size_of::<TxDescriptor>() * index)
            .try_into_u32()
            .unwrap();

        regs.DMA.DMAC0TXDTPR.set(ring_pos);

        Ok(())
    }

    async fn softirq(&self) {
        let inner = self.inner.get();
        let dma = self.dma.get();
        let iface = *self.iface_id.get();

        loop {
            let events = self.softirq_events.wait().await;

            if events & Self::SOFTIRQ_GMII_STATUS != 0 {
                let link_status = inner.link_state();
                log::info!("stmmac: link is {link_status}");
            }

            if events & Self::SOFTIRQ_TX_STATUS != 0 {
                inner.tx_ring.lock().consume().ok();
            }

            if events & Self::SOFTIRQ_RX_STATUS != 0 {
                let mut rx_ring = inner.rx_ring.lock();
                rx_ring
                    .consume(dma.as_ref(), |packet| {
                        let packet = RxPacket::new(packet, 0, iface);
                        ygg_driver_net_core::receive_packet(packet).ok();
                    })
                    .ok();
            }
        }
    }
}

impl InterruptHandler for Stmmac {
    fn handle_irq(self: Arc<Self>, _vector: IrqVector) -> bool {
        let inner = self.inner.get();
        let regs = inner.regs.lock();
        let mac_isr = regs.MAC.MACISR.extract();
        let dma0_sr = regs.DMA.DMAC0SR.extract();

        if mac_isr.matches_all(MACISR::RGSMIIIS::SET) {
            let _ = regs.MAC.MACPHYCSR.get();
            self.softirq_events.signal(Self::SOFTIRQ_GMII_STATUS);
        }

        if dma0_sr.matches_all(DMACiSR::TI::SET) {
            regs.DMA.DMAC0SR.modify(DMACiSR::TI::SET);
            self.softirq_events.signal(Self::SOFTIRQ_TX_STATUS);
        }
        if dma0_sr.matches_all(DMACiSR::RI::SET) {
            regs.DMA.DMAC0SR.modify(DMACiSR::RI::SET);
            self.softirq_events.signal(Self::SOFTIRQ_RX_STATUS);
        }

        true
    }
}

impl Device for Stmmac {
    unsafe fn init(self: Arc<Self>, cx: DeviceInitContext) -> Result<(), Error> {
        // TODO this is a dirty fix for the device tree double initialization problem
        if self.dma.is_initialized() {
            return Err(Error::AlreadyExists);
        }
        let dma = self.dma.init(cx.dma_allocator.clone());

        self.irq.register(self.clone())?;

        let tx_ring_capacity = 32;
        let rx_ring_capacity = 32;

        log::info!("stmmac: setup {:#x}, mac={}", self.base, self.mac);

        self.clk_stmmaceth.enable()?;
        self.clk_pclk.enable()?;
        self.clk_ptp_ref.enable()?;
        self.clk_gtx.enable()?;
        self.clk_tx.enable()?;

        self.rst_ahb.deassert()?;

        self.rst_stmmaceth.assert()?;
        psleep(Duration::from_millis(1));
        self.rst_stmmaceth.deassert()?;

        psleep(Duration::from_millis(20));

        let regs = unsafe { DeviceMemoryIo::<Regs>::map(self.base, Default::default()) }?;

        // DMA initialization

        // Perform a software reset
        regs.DMA.DMAMR.modify(DMAMR::SWR::SET);
        pwait(
            Duration::from_millis(100),
            Duration::from_millis(10),
            || regs.DMA.DMAMR.matches_all(DMAMR::SWR::CLEAR),
        )?;

        // Setup DMA parameters
        // TODO get these params from device tree
        regs.DMA.DMASBMR.write(
            DMASBMR::FB::SET
                + DMASBMR::BLEN256::SET
                + DMASBMR::BLEN128::SET
                + DMASBMR::BLEN64::SET
                + DMASBMR::BLEN32::SET
                + DMASBMR::EN_LPI::SET
                + DMASBMR::WR_OSR_LMT.val(3)
                + DMASBMR::RD_OSR_LMT.val(3),
        );

        // Setup DMA Tx/Rx rings
        let tx_ring = TxRing::with_capacity(dma.as_ref(), tx_ring_capacity)?;
        let rx_ring = RxRing::with_capacity(dma.as_ref(), rx_ring_capacity)?;
        let tx_ring_base = tx_ring.buffer_base().try_into_u32().unwrap();
        let rx_ring_base = rx_ring.buffer_base().try_into_u32().unwrap();

        regs.DMA.DMAC0TXRLR.set(tx_ring_capacity as u32 - 1);
        regs.DMA.DMAC0TXDLAR.set(tx_ring_base);
        regs.DMA.DMAC0TXDTPR.set(tx_ring_base);

        regs.DMA.DMAC0RXRLR.set(rx_ring_capacity as u32 - 1);
        regs.DMA.DMAC0RXDLAR.set(rx_ring_base);
        regs.DMA
            .DMAC0RXDTPR
            .set(rx_ring_base + (rx_ring_capacity * size_of::<RxDescriptor>()) as u32);

        // Setup DMA maximum segmen size, Rx buffer size + max burst len
        regs.DMA.DMAC0CR.write(DMACiCR::PBLX8::CLEAR);
        regs.DMA
            .DMAC0TXCR
            .write(DMACiTXCR::TXPBL.val(32) + DMACiTXCR::OSF::SET);
        regs.DMA
            .DMAC0RXCR
            .write(DMAC0RXCR::RBSZ.val(4096) + DMAC0RXCR::RXPBL.val(32));

        // Enable DMA interrupts
        // TODO enable abnormal interrupts to handle errors properly
        regs.DMA.DMAC0IER.write(
            DMACiIER::NIE::SET
                + DMACiIER::RIE::SET
                + DMACiIER::TIE::SET
                + DMACiIER::RBUE::SET
                + DMACiIER::RSE::SET,
        );

        // Start Tx/Rx DMAs
        regs.DMA.DMAC0TXCR.modify(DMACiTXCR::ST::SET);
        regs.DMA.DMAC0RXCR.modify(DMAC0RXCR::SR::SET);

        // MTL initialization

        regs.MTL
            .MTLOMR
            .write(MTLOMR::RAA::AlgStrictPrio + MTLOMR::SCHALG::AlgStrictPrio);

        // TODO get TQS, RQS from device tree
        regs.MTL
            .MTLTXQ0OMR
            .write(MTLTXQiOMR::TSF::SET + MTLTXQiOMR::TQS.val(7) + MTLTXQiOMR::TXQEN::Enable);
        regs.MTL.MTLRXQ0OMR.write(
            MTLRXQiOMR::RQS.val(7)
                + MTLRXQiOMR::DIS_TCP_EF::SET
                + MTLRXQiOMR::RSF::SET
                + MTLRXQiOMR::FEP::SET
                + MTLRXQiOMR::FUP::SET,
        );

        // MAC initialization

        // Setup the MAC address
        let mac_bytes: [u8; 6] = self.mac.into();
        regs.MAC.MACA0HR.write(
            MACAiHR::AE::SET
                + MACAiHR::ADDR4.val(mac_bytes[4] as u32)
                + MACAiHR::ADDR5.val(mac_bytes[5] as u32),
        );
        regs.MAC.MACA0LR.write(
            MACAiLR::ADDR0.val(mac_bytes[0] as u32)
                + MACAiLR::ADDR1.val(mac_bytes[1] as u32)
                + MACAiLR::ADDR2.val(mac_bytes[2] as u32)
                + MACAiLR::ADDR3.val(mac_bytes[3] as u32),
        );

        // TODO proper filtering, don't use promiscuous mode all the time
        regs.MAC.MACPFR.write(MACPFR::PR::SET + MACPFR::RA::SET);

        // Setup Tx/Rx flow control, enable Rx queue 0
        regs.MAC.MACQ0TXFCR.write(MACQ0TXFCR::TFE::SET);
        regs.MAC.MACRXFCR.write(MACRXFCR::RFE::SET);
        regs.MAC.MACRXQC0R.write(MACRXQC0R::RXQ0EN::Enable);

        // Setup MAC-level interrupts
        regs.MAC
            .MACIER
            .write(MACIER::TXSTSIE::SET + MACIER::RXSTSIE::SET + MACIER::RGSMIIIE::SET);

        // Setup link information
        regs.MAC
            .MACCR
            .modify(MACCR::PS::Ps1000Mbps + MACCR::DM::FullDuplex);

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

        // Apply PHY settings
        regs.MAC
            .MACPHYCSR
            .modify(MACPHYCSR::LUD::SET + MACPHYCSR::TC::SET);

        // Setup the PHY
        // TODO autodiscover phy
        let phy = PhyAccess::new(&*regs, 0x00);
        let (id0, id1) = phy.id()?;
        log::info!("stmmac: PHY {id0:04x}:{id1:04x}");
        phy.reset(Duration::from_millis(100))?;
        phy.setup_link(true, Some(GBESR::empty()))?;

        self.inner.init(Inner {
            regs: IrqSafeSpinlock::new(regs),

            tx_ring: IrqSafeSpinlock::new(tx_ring),
            rx_ring: IrqSafeSpinlock::new(rx_ring),
        });

        let iface =
            ygg_driver_net_core::register_interface(NetworkInterfaceType::Ethernet, self.clone());
        self.iface_id.init(iface.id());

        self.irq.enable()?;

        let p = self.clone();
        runtime::spawn(async move { p.softirq().await })?;

        Ok(())
    }

    fn display_name(&self) -> &str {
        "Synopsys DesignWare Ethernet MAC"
    }
}

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

    fn transmit_buffer(&self, buffer: DmaBuffer<[u8]>) -> Result<(), Error> {
        self.start_xmit(buffer)
    }

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

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

    fn link_state(&self) -> LinkState {
        self.inner.get().link_state()
    }
}

device_tree_driver! {
    compatible: ["starfive,jh7110-dwmac"],
    driver: {
        fn probe(&self, node: &Arc<Node>, context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            let base = node.map_base(context, 0)?;

            let clk_stmmaceth = node.named_clock("stmmaceth")?;
            let clk_pclk = node.named_clock("pclk")?;
            let clk_ptp_ref = node.named_clock("ptp_ref")?;
            let clk_tx = node.named_clock("tx")?;
            let clk_gtx = node.named_clock("gtx")?;

            let rst_stmmaceth = node.named_reset("stmmaceth")?;
            let rst_ahb = node.named_reset("ahb")?;

            let mac = read_mac_address(node)?;
            // TODO named_interrupt()
            let irq = node.interrupt(0)?;

            context.sequence = Some(InitSequence::Late);

            Some(Arc::new(Stmmac {
                base,
                mac,
                irq,

                clk_stmmaceth,
                clk_pclk,
                clk_ptp_ref,
                clk_tx,
                clk_gtx,
                rst_stmmaceth,
                rst_ahb,

                softirq_events: BitmapEvent::new(AtomicWaker::new()),

                dma: OneTimeInit::new(),
                inner: OneTimeInit::new(),
                iface_id: OneTimeInit::new(),
            }))
        }
    }
}
