use alloc::sync::Arc;
use device_api::{
    clock::{ClockHandle, Hertz},
    device::{Device, DeviceInitContext},
    i2c::{I2CAddress, I2CController},
    interrupt::{InterruptHandler, IrqHandle, IrqVector},
};
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::{OneTimeInit, sync::spin_rwlock::IrqSafeRwLock};
use tock_registers::{
    interfaces::{ReadWriteable, Readable, Writeable},
    register_bitfields, register_structs,
    registers::ReadWrite,
};
use yggdrasil_abi::io::device::i2c::I2CMasterInfo;

register_bitfields! {
    u32,
    C [
        I2CEN   OFFSET(15)  NUMBITS(1)  [],
        INTR    OFFSET(10)  NUMBITS(1)  [],
        INTT    OFFSET(9)   NUMBITS(1)  [],
        INTD    OFFSET(8)   NUMBITS(1)  [],
        ST      OFFSET(7)   NUMBITS(1)  [],
        CLEAR   OFFSET(4)   NUMBITS(2)  [],
        READ    OFFSET(0)   NUMBITS(1)  [],
    ],
    S [
        CLKT    OFFSET(9)   NUMBITS(1)  [],
        ERR     OFFSET(8)   NUMBITS(1)  [],
        RXF     OFFSET(7)   NUMBITS(1)  [],
        TXE     OFFSET(6)   NUMBITS(1)  [],
        RXD     OFFSET(5)   NUMBITS(1)  [],
        TXD     OFFSET(4)   NUMBITS(1)  [],
        RXR     OFFSET(3)   NUMBITS(1)  [],
        TXW     OFFSET(2)   NUMBITS(1)  [],
        DONE    OFFSET(1)   NUMBITS(1)  [],
        TA      OFFSET(0)   NUMBITS(1)  [],
    ],
}

register_structs! {
    #[allow(non_snake_case)]
    Regs {
        (0x00 => C: ReadWrite<u32, C::Register>),
        (0x04 => S: ReadWrite<u32, S::Register>),
        (0x08 => DLEN: ReadWrite<u32>),
        (0x0C => A: ReadWrite<u32>),
        (0x10 => FIFO: ReadWrite<u32>),
        (0x14 => DIV: ReadWrite<u32>),
        (0x18 => DEL: ReadWrite<u32>),
        (0x1C => CLKT: ReadWrite<u32>),
        (0x20 => @END),
    }
}

struct I2C {
    name: &'static str,
    clock_frequency: Option<Hertz>,
    irq: IrqHandle,
    clock: Option<ClockHandle>,
    regs: IrqSafeRwLock<DeviceMemoryIo<'static, Regs>>,
    index: OneTimeInit<u32>,
}

impl Regs {
    fn start_transfer(&self, name: &str, buflen: u16, address: I2CAddress, read: bool) {
        log::debug!(
            "{}: start address={}, read={}, len={}",
            name,
            address,
            read,
            buflen
        );

        let address = address.as_8_bit().unwrap();
        let read = if read { C::READ::SET } else { C::READ::CLEAR };

        self.S.write(S::ERR::SET + S::DONE::SET);
        self.DLEN.set(buflen as u32);
        self.C.modify(C::ST::CLEAR + C::I2CEN::CLEAR);
        self.A.set(address as u32);
        self.C.modify(read + C::ST::SET + C::I2CEN::SET);
    }

    fn finish_transfer(&self, name: &str) -> Result<(), Error> {
        log::debug!("{}: finish transfer", name);
        let status = self.S.extract();
        self.C.set(0);
        if status.matches_all(S::ERR::SET) {
            self.S.write(S::DONE::SET + S::ERR::SET);
            return Err(Error::HostUnreachable);
        }
        if status.matches_all(S::DONE::SET) {
            self.S.write(S::DONE::SET);
        }

        Ok(())
    }

    fn write_byte(&self, byte: u8) -> Result<bool, Error> {
        loop {
            let status = self.S.extract();
            if status.matches_all(S::ERR::SET) {
                self.C.write(C::CLEAR.val(1));
                self.S.write(S::ERR::SET + S::DONE::SET);
                // TODO better code
                return Err(Error::HostUnreachable);
            }
            if status.matches_all(S::DONE::SET) {
                self.C.set(0);
                self.S.write(S::DONE::SET);
                return Ok(false);
            }
            if status.matches_all(S::TXD::SET) {
                self.FIFO.set(byte as u32);
                return Ok(true);
            }

            core::hint::spin_loop();
        }
    }

    fn read_byte(&self) -> Result<Option<u8>, Error> {
        loop {
            let status = self.S.extract();
            if status.matches_all(S::ERR::SET) {
                self.C.write(C::CLEAR.val(1));
                self.S.write(S::ERR::SET + S::DONE::SET);
                // TODO better code
                return Err(Error::HostUnreachable);
            }
            if status.matches_all(S::DONE::SET) {
                self.C.set(0);
                self.S.write(S::DONE::SET);
                return Ok(None);
            }
            if status.matches_all(S::RXD::SET) {
                let val = self.FIFO.get() as u8;
                return Ok(Some(val));
            }

            core::hint::spin_loop();
        }
    }
}

impl Device for I2C {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        // TODO
        let _ = &self.clock;

        log::info!("{}: initialize", self.name);
        let regs = self.regs.write();
        regs.C.set(0);

        let index = DEVICE_REGISTRY.i2c.register_bus(self.clone())?;
        self.index.init(index);
        Ok(())
    }

    unsafe fn init_irq(self: Arc<Self>) -> Result<(), Error> {
        // TODO
        let _ = &self.irq;
        // self.irq.register(self.clone())?;
        // self.irq.enable()?;
        // let regs = self.regs.write();
        // regs.C.modify(C::INTD::SET);
        Ok(())
    }

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

impl InterruptHandler for I2C {
    fn handle_irq(self: Arc<Self>, _vector: IrqVector) -> bool {
        todo!()
    }
}

impl I2CController for I2C {
    fn bus_number(&self) -> u32 {
        *self.index.get()
    }

    fn child_number(&self) -> Option<u32> {
        None
    }

    fn set_speed(&self, speed: Hertz) -> Result<Hertz, Error> {
        if let Some(frequency) = self.clock_frequency {
            // Fixed frequency
            if speed == frequency {
                Ok(speed)
            } else {
                Err(Error::InvalidArgument)
            }
        } else {
            todo!("Set i2c speed = {speed}")
        }
    }

    fn capabilities(&self) -> I2CMasterInfo {
        let max_speed_hz = if let Some(frequency) = self.clock_frequency {
            frequency
        } else {
            todo!()
        }
        .0 as u32;
        I2CMasterInfo {
            max_speed_hz,
            supports_10bit_addresses: true,
        }
    }

    fn i2c_write(&self, address: I2CAddress, buffer: &[u8]) -> Result<usize, Error> {
        let buflen: u16 = buffer
            .len()
            .try_into()
            .map_err(|_| Error::InvalidArgument)?;
        let regs = self.regs.write();
        // TODO DMA/interrupts
        regs.start_transfer(self.name, buflen, address, false);
        let mut bytes_written = 0;
        // let mut done = false;
        for &byte in buffer {
            if regs.write_byte(byte)? {
                bytes_written += 1;
            } else {
                break;
            }
        }
        regs.finish_transfer(self.name)?;
        Ok(bytes_written)
    }

    fn i2c_read(&self, address: I2CAddress, buffer: &mut [u8]) -> Result<usize, Error> {
        let buflen: u16 = buffer
            .len()
            .try_into()
            .map_err(|_| Error::InvalidArgument)?;
        let regs = self.regs.write();
        // TODO DMA/interrupts
        regs.start_transfer(self.name, buflen, address, true);
        let mut bytes_read = 0;
        // let mut done = false;
        for byte in buffer {
            if let Some(val) = regs.read_byte()? {
                *byte = val;
                bytes_read += 1;
            } else {
                break;
            }
        }
        regs.finish_transfer(self.name)?;
        Ok(bytes_read)
    }
}

device_tree_driver! {
    compatible: ["brcm,bcm2711-i2c", "brcm,bcm2835-i2c"],
    driver: {
        fn probe(&self, node: &Arc<Node>, context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            let clock_frequency = node.prop_usize("clock-frequency").map(|p| Hertz(p as u64));
            let base = node.map_base(context, 0)?;
            let name = node.name().unwrap_or("bcm2835-i2c");
            let irq = node.interrupt(0)?;
            let clock = node.clock(0);

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

            let i2c = Arc::new(I2C {
                name,
                clock_frequency,
                irq,
                clock,
                index: OneTimeInit::new(),
                regs: IrqSafeRwLock::new(regs)
            });

            node.make_i2c_controller(i2c.clone());

            Some(i2c)
        }
    }
}
