use core::iter;

use alloc::{boxed::Box, sync::Arc};
use async_trait::async_trait;
use device_api::{
    clock::{ClockHandle, Hertz},
    device::{Device, DeviceInitContext},
    interrupt::IrqHandle,
    spi::SpiController,
};
use device_tree::driver::{Node, ProbeContext, device_tree_driver};
use libk::{device::manager::DEVICE_REGISTRY, error::Error};
use libk_mm::device::DeviceMemoryIo;
use libk_util::sync::IrqSafeSpinlock;
use tock_registers::{
    interfaces::{Readable, Writeable},
    register_bitfields, register_structs,
    registers::ReadWrite,
};
use yggdrasil_abi::io::device::spi::{SpiClockPhase, SpiClockPolarity, SpiConfig};

register_bitfields! {
    u32,
    CS [
        LEN_LONG    OFFSET(25)  NUMBITS(1)  [],
        DMA_LEN     OFFSET(24)  NUMBITS(1)  [],
        CSPOL2      OFFSET(23)  NUMBITS(1)  [],
        CSPOL1      OFFSET(22)  NUMBITS(1)  [],
        CSPOL0      OFFSET(21)  NUMBITS(1)  [],
        RXF         OFFSET(20)  NUMBITS(1)  [],
        RXR         OFFSET(19)  NUMBITS(1)  [],
        TXD         OFFSET(18)  NUMBITS(1)  [],
        RXD         OFFSET(17)  NUMBITS(1)  [],
        DONE        OFFSET(16)  NUMBITS(1)  [],
        LEN         OFFSET(13)  NUMBITS(1)  [],
        REN         OFFSET(12)  NUMBITS(1)  [],
        ADCS        OFFSET(11)  NUMBITS(1)  [],
        INTR        OFFSET(10)  NUMBITS(1)  [],
        INTD        OFFSET(9)   NUMBITS(1)  [],
        DMAEN       OFFSET(8)   NUMBITS(1)  [],
        TA          OFFSET(7)   NUMBITS(1)  [],
        CSPOL       OFFSET(6)   NUMBITS(1)  [],
        CLEAR       OFFSET(4)   NUMBITS(2)  [],
        CPOL        OFFSET(3)   NUMBITS(1)  [
            IdleLow = 0,
            IdleHigh = 1
        ],
        CPHA        OFFSET(2)   NUMBITS(1)  [
            CaptureOnFirstTransition = 0,
            CaptureOnSecondTransition = 1
        ],
        CS          OFFSET(0)   NUMBITS(2)  [],
    ]
}

register_structs! {
    #[allow(non_snake_case)]
    Regs {
        (0x00 => CS: ReadWrite<u32, CS::Register>),
        (0x04 => FIFO: ReadWrite<u32>),
        (0x08 => CLK: ReadWrite<u32>),
        (0x0C => DLEN: ReadWrite<u32>),
        (0x10 => LTOH: ReadWrite<u32>),
        (0x14 => DC: ReadWrite<u32>),
        (0x18 => @END),
    }
}

pub struct Spi {
    name: &'static str,
    clock: ClockHandle,
    irq: IrqHandle,
    regs: IrqSafeSpinlock<DeviceMemoryIo<'static, Regs>>,
}

impl Device for Spi {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        let _ = &self.irq;
        let regs = self.regs.lock();
        regs.CS.set(0);

        DEVICE_REGISTRY.spi.register_bus(self.clone()).ok();
        Ok(())
    }

    fn display_name(&self) -> &str {
        self.name
    }
}

impl Regs {
    fn start_transfer(&self, core_clock: Hertz, config: &SpiConfig) -> Result<(), Error> {
        let want_clock = Hertz::from(config.frequency.unwrap_or(100000));
        let divider = Hertz::divider(core_clock, want_clock).ok_or(Error::InvalidArgument)?;
        let divider: u16 = divider.try_into().map_err(|_| Error::InvalidArgument)?;
        let divider = divider.wrapping_add(1) & !1;

        let mut cs = CS::REN::SET + CS::TA::SET;
        if let Some(cs_index) = config.chip_select_index
            && cs_index < 3
        {
            cs += CS::CS.val(cs_index);
        }
        if let Some(cpol) = config.clock_polarity {
            match cpol {
                SpiClockPolarity::IdleLow => cs += CS::CPOL::IdleLow,
                SpiClockPolarity::IdleHigh => cs += CS::CPOL::IdleHigh,
            }
        }
        if let Some(cpha) = config.clock_phase {
            match cpha {
                SpiClockPhase::CaptureOnFirstTransition => cs += CS::CPHA::CaptureOnFirstTransition,
                SpiClockPhase::CaptureOnSecondTransition => {
                    cs += CS::CPHA::CaptureOnSecondTransition
                }
            }
        }

        self.CLK.set(divider as u32);
        self.CS.write(cs);

        Ok(())
    }

    fn tx_rx_byte(&self, tx: u8) -> Result<Option<u8>, Error> {
        // Wait for tx
        loop {
            let status = self.CS.extract();
            if status.matches_all(CS::TXD::SET) {
                break;
            }
            if status.matches_all(CS::DONE::SET) {
                return Ok(None);
            }
        }
        self.FIFO.set(tx as u32);
        // Wait for rx
        loop {
            let status = self.CS.extract();
            if status.matches_all(CS::RXD::SET) {
                break;
            }
            if status.matches_all(CS::DONE::SET) {
                return Ok(None);
            }
        }
        Ok(Some(self.FIFO.get() as u8))
    }
}

#[async_trait]
impl SpiController for Spi {
    async fn spi_transfer(
        &self,
        tx_buffer: &[u8],
        rx_buffer: &mut [u8],
        config: &SpiConfig,
    ) -> Result<usize, Error> {
        // TODO interrupt/DMA
        log::debug!("{}: spi_transfer {} bytes", self.name, rx_buffer.len());
        let core_clock = self.clock.rate()?;
        let regs = self.regs.lock();
        regs.start_transfer(core_clock, config)?;

        let mut len = 0;
        for (&tx, rx) in iter::zip(tx_buffer, rx_buffer) {
            match regs.tx_rx_byte(tx)? {
                Some(b) => {
                    *rx = b;
                    len += 1;
                }
                None => {
                    break;
                }
            }
        }
        Ok(len)
    }
}

device_tree_driver! {
    compatible: ["brcm,bcm2835-spi"],
    driver: {
        fn probe(&self, node: &Arc<Node>, context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            let name = node.name().unwrap_or("bcm2835-spi");
            let base = node.map_base(context, 0)?;
            let irq = node.interrupt(0)?;
            let clock = node.clock(0)?;

            let regs = unsafe { DeviceMemoryIo::map(base, Default::default()) }.ok()?;

            let spi = Arc::new(Spi {
                regs: IrqSafeSpinlock::new(regs),
                name,
                clock,
                irq
            });

            Some(spi)
        }
    }
}
