use alloc::{boxed::Box, sync::Arc};
use async_trait::async_trait;
use device_api::{
    clock::{ClockHandle, Hertz},
    device::{Device, DeviceInitContext},
    i2c::{I2CAddress, I2CController, I2CMessage, I2CTransfer},
    interrupt::{InterruptHandler, IrqHandle, IrqVector},
};
use device_tree::driver::{Node, ProbeContext, device_tree_driver};
use futures_util::task::AtomicWaker;
use libk::{device::manager::DEVICE_REGISTRY, error::Error};
use libk_mm::device::DeviceMemoryIo;
use libk_util::{OneTimeInit, event::BitmapEvent, sync::spin_rwlock::IrqSafeRwLock};
use tock_registers::{
    LocalRegisterCopy,
    fields::FieldValue,
    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 I2CState {
    status: BitmapEvent<AtomicWaker>,
}

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

impl I2CState {
    fn new() -> Self {
        Self {
            status: BitmapEvent::new(AtomicWaker::new()),
        }
    }
}

// All of these are thread-unsafe, but I2C subsystem locks the peripheral by process
impl I2C {
    async fn wait_for_event(&self) -> LocalRegisterCopy<u32, S::Register> {
        let event = self.state.status.wait().await;
        LocalRegisterCopy::new(event as u32)
    }
}

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> {
        self.irq.register(self.clone())?;
        self.irq.enable()?;
        Ok(())
    }

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

impl InterruptHandler for I2C {
    fn handle_irq(self: Arc<Self>, _vector: IrqVector) -> bool {
        let regs = self.regs.write();
        let event = regs.S.extract();
        // log::info!("I2C irq {:#x}", event.get());
        let mut clear = FieldValue::none();
        if event.matches_all(S::RXR::SET) {
            clear += C::INTR::CLEAR;
        }
        if event.matches_all(S::TXW::SET) {
            clear += C::INTT::CLEAR;
        }
        if event.matches_all(S::DONE::SET) {
            clear += C::INTD::CLEAR;
        }
        regs.C.modify(clear);
        self.state.status.signal(event.get() as u64);
        true
    }
}

#[async_trait]
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,
        }
    }

    async fn i2c_transfer(
        &self,
        address: I2CAddress,
        transfer: &mut I2CTransfer,
    ) -> Result<usize, Error> {
        // TODO validate transfer
        let address = address.as_8_bit().ok_or(Error::NotImplemented)?;
        let repeat_start = transfer.repeat_start;
        let mut nbytes = 0;
        let mut error = None;
        for (i, message) in transfer.messages.iter_mut().enumerate() {
            if i == 0 || repeat_start {
                let (read, len) = match message {
                    I2CMessage::Read(buffer) => (true, buffer.len()),
                    I2CMessage::Write(buffer) => (false, buffer.len()),
                };
                // log::info!("{}: START {:#x} read={}", self.name, address, read);
                let read = if read { C::READ::SET } else { C::READ::CLEAR };
                let regs = self.regs.write();
                regs.S.write(S::ERR::SET + S::DONE::SET);
                regs.DLEN.set(len as u32);
                regs.C.modify(C::ST::CLEAR);
                regs.A.set(address as u32);
                regs.C.modify(
                    read + C::ST::SET + C::I2CEN::SET + C::INTD::SET + C::INTR::SET + C::INTT::SET,
                );
            }

            let mut pos = 0;
            let status = loop {
                let s = self.wait_for_event().await;
                let regs = self.regs.write();
                if s.matches_all(S::ERR::SET) || s.matches_all(S::DONE::SET) {
                    break s;
                }
                match message {
                    I2CMessage::Write(buffer) => {
                        if s.matches_all(S::TXW::SET) {
                            while regs.S.matches_all(S::TXD::SET) && pos < buffer.len() {
                                let byte = buffer[pos];
                                // log::info!("{}: tx {:#x}", self.name, byte);
                                regs.FIFO.set(byte as u32);
                                pos += 1;
                                nbytes += 1;
                            }
                        }
                        regs.C.modify(C::INTT::SET);
                    }
                    I2CMessage::Read(buffer) => {
                        if s.matches_all(S::RXR::SET) {
                            while regs.S.matches_all(S::RXD::SET) && pos < buffer.len() {
                                let byte = regs.FIFO.get() as u8;
                                // log::info!("{}: rx {:#x}", self.name, byte);
                                buffer[pos] = byte;
                                pos += 1;
                                nbytes += 1;
                            }
                        }
                        regs.C.modify(C::INTR::SET);
                    }
                }
            };

            // Finalize message
            let regs = self.regs.write();
            if status.matches_all(S::ERR::SET) {
                regs.S.write(S::DONE::SET + S::ERR::SET);
                error = Some(Error::HostUnreachable);
                break;
            } else if status.matches_all(S::DONE::SET) {
                regs.S.write(S::DONE::SET);
            } else {
                log::error!("{}: transfer finished without DONE or ERR set", self.name);
                log::error!("{}: S = {:#x}", self.name, status.get());
                todo!();
            }
        }

        // Finish transfer
        // log::info!(
        //     "{}: finish xfer, error={:?}, nbytes={}",
        //     self.name,
        //     error,
        //     nbytes
        // );
        let regs = self.regs.write();
        regs.C.set(0);

        match error {
            Some(error) => Err(error),
            None => Ok(nbytes),
        }
    }
}

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),
                state: I2CState::new()
            });

            node.make_i2c_controller(i2c.clone());

            Some(i2c)
        }
    }
}
