use alloc::{boxed::Box, sync::Arc};
use async_trait::async_trait;
use device_api::{
    clock::{ClockHandle, Hertz, IntoHertz},
    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::{block, device::manager::DEVICE_REGISTRY, error::Error};
use libk_mm::device::DeviceMemoryIo;
use libk_util::{OneTimeInit, event::BitmapEvent, sync::IrqSafeSpinlock};
use tock_registers::{
    LocalRegisterCopy,
    interfaces::{ReadWriteable, Readable, Writeable},
    register_bitfields, register_structs,
    registers::{Aliased, ReadWrite},
};
use yggdrasil_abi::io::device::i2c::I2CMasterInfo;

const CMD_START: u32 = 0x91;
const CMD_STOP: u32 = 0x41;
const CMD_WRITE: u32 = 0x11;
const CMD_READ_ACK: u32 = 0x21;
const CMD_READ_NACK: u32 = 0x29;
const CMD_IACK: u32 = 0x01;

register_bitfields! {
    u32,
    CTR [
        EN      OFFSET(7)   NUMBITS(1)  [],
        IEN     OFFSET(6)   NUMBITS(1)  [],
    ],
    SR [
        RXACK   OFFSET(7)   NUMBITS(1)  [],
        BUSY    OFFSET(6)   NUMBITS(1)  [],
        AL      OFFSET(5)   NUMBITS(1)  [],
        TIP     OFFSET(1)   NUMBITS(1)  [],
        IF      OFFSET(0)   NUMBITS(1)  [],
    ],
}

register_structs! {
    #[allow(non_snake_case)]
    Regs {
        (0x000 => PRER_LO: ReadWrite<u32>),
        (0x004 => PRER_HI: ReadWrite<u32>),
        (0x008 => CTR: ReadWrite<u32, CTR::Register>),
        (0x00C => TXR_RXR: ReadWrite<u32>),
        (0x010 => CR_SR: Aliased<u32, SR::Register, ()>),
        (0x014 => _0),
        (0x020 => @END),
    }
}

struct AbortGuard<'a> {
    i2c: &'a I2C,
    abort: bool,
}

impl<'a> AbortGuard<'a> {
    pub fn new(i2c: &'a I2C) -> Self {
        Self { i2c, abort: true }
    }

    pub fn consume(mut self) {
        self.abort = false;
    }
}

impl Drop for AbortGuard<'_> {
    fn drop(&mut self) {
        if self.abort {
            block!(self.i2c.finish_transfer(Err(Error::Interrupted)).await).ok();
        }
    }
}

pub struct I2C {
    name: &'static str,
    clock: ClockHandle,
    irq: IrqHandle,
    regs: IrqSafeSpinlock<DeviceMemoryIo<'static, Regs>>,
    index: OneTimeInit<u32>,
    sr: BitmapEvent<AtomicWaker>,
}

impl Regs {
    fn set_clock_rate(&self, input_rate: Hertz, clock_rate: Hertz) -> Result<Hertz, Error> {
        let divider = Hertz::divider(input_rate / 5, clock_rate).ok_or(Error::InvalidArgument)?;
        if divider == 0 || divider > 0xFFFF {
            return Err(Error::InvalidArgument);
        }
        log::info!("input_rate = {input_rate}, desired = {clock_rate}, divider = {divider}");
        let divider = divider - 1;
        self.PRER_LO.set(divider & 0xFF);
        self.PRER_HI.set(divider >> 8);
        let real_rate = input_rate / (5 * (divider + 1));
        Ok(real_rate)
    }
}

impl I2C {
    async fn wait_for_event(&self) -> LocalRegisterCopy<u32, SR::Register> {
        let value = self.sr.wait().await;
        LocalRegisterCopy::new(value as u32)
    }

    async fn finish_transfer(&self, status: Result<usize, Error>) -> Result<usize, Error> {
        if status == Err(Error::Interrupted) {
            log::warn!("{}: transfer aborted", self.name);

            let regs = self.regs.lock();
            let stopped = {
                if regs.CR_SR.matches_all(SR::TIP::CLEAR) {
                    true
                } else {
                    false
                }
            };

            if !stopped {
                regs.CTR.modify(CTR::IEN::SET);
                drop(regs);

                loop {
                    let sr = self.wait_for_event().await;
                    if sr.matches_all(SR::TIP::CLEAR) {
                        break;
                    }
                    let regs = self.regs.lock();
                    regs.CTR.modify(CTR::IEN::SET); // Reenable IRQs
                }
            }
        }
        // Generate stop condition
        {
            let regs = self.regs.lock();
            regs.CR_SR.set(CMD_STOP);
            regs.CTR.modify(CTR::IEN::SET); // Reenable IRQs
        }

        loop {
            let sr = self.wait_for_event().await;
            if sr.matches_all(SR::BUSY::CLEAR) {
                break;
            }
            let regs = self.regs.lock();
            regs.CTR.modify(CTR::IEN::SET); // Reenable IRQs
        }

        let regs = self.regs.lock();
        regs.CTR.modify(CTR::IEN::CLEAR); // Mask IRQs
        regs.CR_SR.set(CMD_IACK);

        match status {
            Ok(count) => Ok(count),
            Err(err) => Err(err),
        }
    }
}

impl Device for I2C {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        let input_clk = self.clock.rate()?;
        let regs = self.regs.lock();
        regs.CTR.set(0);
        regs.set_clock_rate(input_clk, 100u64.khz())?;
        regs.CTR.write(CTR::EN::SET);
        regs.CR_SR.set(CMD_IACK);
        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.lock();
        let status = regs.CR_SR.extract();
        self.sr.signal(status.get() as u64);
        regs.CTR.modify(CTR::IEN::CLEAR);
        true
    }
}

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

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

    fn capabilities(&self) -> I2CMasterInfo {
        let scl_max = self.clock.rate().expect("Couldn't get input clock rate") / 5;
        I2CMasterInfo {
            max_speed_hz: scl_max.0 as u32,
            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;
        let mut abort = None;
        for (i, message) in transfer.messages.iter_mut().enumerate() {
            let (read, buffer_len) = match message {
                I2CMessage::Read(buffer) => (true, buffer.len()),
                I2CMessage::Write(buffer) => (false, buffer.len()),
            };

            if i == 0 || repeat_start {
                //log::info!("{}: START {:#x} read={}", self.name, address, read);
                let read_bit = if read { 1 << 0 } else { 0 << 0 };
                let regs = self.regs.lock();
                regs.TXR_RXR.set(((address as u32) << 1) | read_bit);
                regs.CR_SR.set(CMD_START);
                regs.CTR.modify(CTR::IEN::SET);
                if abort.is_none() {
                    abort = Some(AbortGuard::new(self));
                }
            }

            let mut pos = 0;
            let mut start = true;
            loop {
                if pos == buffer_len {
                    break;
                }
                let sr = self.wait_for_event().await;
                let regs = self.regs.lock();
                if sr.matches_all(SR::AL::SET) {
                    error = Some(Error::HostUnreachable);
                    break;
                }

                match message {
                    I2CMessage::Read(buffer) => {
                        if pos == 0 && sr.matches_all(SR::RXACK::SET) {
                            error = Some(Error::HostUnreachable);
                            break;
                        }
                        if start {
                            regs.CR_SR.set(CMD_READ_ACK);
                            start = false;
                        } else if sr.matches_all(SR::TIP::CLEAR) {
                            let byte = regs.TXR_RXR.get() as u8;
                            //log::info!("{}: rx {:#x}", self.name, byte);
                            buffer[pos] = byte;
                            pos += 1;
                            nbytes += 1;
                            if pos >= buffer.len() {
                                regs.CR_SR.set(CMD_READ_NACK);
                            } else {
                                regs.CR_SR.set(CMD_READ_ACK);
                            }
                        }
                    }
                    I2CMessage::Write(buffer) => {
                        if sr.matches_all(SR::RXACK::SET) {
                            error = Some(Error::HostUnreachable);
                            break;
                        }
                        regs.TXR_RXR.set(buffer[pos] as u32);
                        //log::info!("{}: tx {:#x}", self.name, buffer[pos]);
                        pos += 1;
                        nbytes += 1;
                        regs.CR_SR.set(CMD_WRITE);
                    }
                }
                regs.CTR.modify(CTR::IEN::SET);
            }

            if error.is_some() {
                break;
            }

            {
                let regs = self.regs.lock();
                if regs.CR_SR.matches_all(SR::TIP::CLEAR) {
                    regs.CTR.modify(CTR::IEN::SET); // Reenable IRQs
                    continue;
                }

                regs.CTR.modify(CTR::IEN::SET); // Reenable IRQs
            }

            loop {
                let sr = self.wait_for_event().await;
                if sr.matches_all(SR::TIP::CLEAR) {
                    break;
                }
                let regs = self.regs.lock();
                regs.CTR.modify(CTR::IEN::SET); // Reenable IRQs
            }
        }

        // Generate stop condition
        {
            //log::info!("{}: STOP", self.name);
            let regs = self.regs.lock();
            regs.CR_SR.set(CMD_STOP);
            regs.CTR.modify(CTR::IEN::SET); // Reenable IRQs
        }

        loop {
            let sr = self.wait_for_event().await;
            if sr.matches_all(SR::BUSY::CLEAR) {
                break;
            }
            let regs = self.regs.lock();
            regs.CTR.modify(CTR::IEN::SET); // Reenable IRQs
        }

        if let Some(abort) = abort.take() {
            abort.consume();
        }

        let regs = self.regs.lock();
        regs.CTR.modify(CTR::IEN::CLEAR); // Mask IRQs
        regs.CR_SR.set(CMD_IACK);

        // log::info!(
        //     "{}: finish xfer, error={:?}, nbytes={}",
        //     self.name,
        //     error,
        //     nbytes
        // );
        match error {
            Some(err) => Err(err),
            None => Ok(nbytes),
        }
    }
    // async fn i2c_read(&self, address: I2CAddress, buffer: &mut [u8]) -> Result<usize, Error> {
    //     self.start_transfer(address, true)?;
    //     let mut pos = 0;
    //     let status = loop {
    //         if pos >= buffer.len() {
    //             break Ok(pos);
    //         }
    //         let sr = self.wait_for_event().await;
    //         let regs = self.regs.lock();
    //         if sr.matches_all(SR::RXACK::SET) || sr.matches_all(SR::AL::SET) {
    //             break Err(sr);
    //         }
    //     };
    //     self.finish_transfer(status).await
    // }

    // async fn i2c_write(&self, address: I2CAddress, buffer: &[u8]) -> Result<usize, Error> {
    //     self.start_transfer(address, false)?;
    //     let mut pos = 0;
    //     let status = loop {
    //         if pos >= buffer.len() {
    //             break Ok(pos);
    //         }
    //         let sr = self.wait_for_event().await;
    //         let regs = self.regs.lock();
    //         if sr.matches_all(SR::RXACK::SET) || sr.matches_all(SR::AL::SET) {
    //             break Err(sr);
    //         }
    //     };
    //     self.finish_transfer(status).await
    // }

    // fn i2c_read(&self, address: I2CAddress, buffer: &mut [u8]) -> Result<usize, Error> {
    //     let address = address.as_8_bit().ok_or(Error::NotImplemented)?;
    //     let regs = self.regs.lock();
    //     log::info!(":::: start {:#x} read", address);
    //     regs.poll_wait(SR::BUSY::SET);
    //     regs.TXR_RXR.set(((address as u32) << 1) | 1);
    //     regs.CR_SR.set(CMD_START);

    //     let mut read = 0;
    //     let mut error = None;
    //     let mut start = true;
    //     loop {
    //         let sr = regs.poll_wait(SR::TIP::SET);
    //         if read == buffer.len() {
    //             break;
    //         }
    //         if sr.matches_all(SR::AL::SET) {
    //             error = Some(Error::WouldBlock);
    //             break;
    //         }
    //         if sr.matches_all(SR::RXACK::SET) {
    //             error = Some(Error::DoesNotExist);
    //             break;
    //         }

    //         if !start {
    //             let val = regs.TXR_RXR.get() as u8;
    //             buffer[read] = val;
    //             read += 1;
    //             if read >= buffer.len() {
    //                 log::info!(":::: read_ack {:02x}", val);
    //                 regs.CR_SR.set(CMD_READ_NACK);
    //             } else {
    //                 log::info!(":::: read {:02x}", val);
    //                 regs.CR_SR.set(CMD_READ_ACK);
    //             }
    //         } else {
    //             log::info!(":::: begin read");
    //             regs.CR_SR.set(CMD_READ);
    //             start = false;
    //         }
    //     }

    //     log::info!(":::: stop {:#x}", address);
    //     // Generate stop condition
    //     regs.CR_SR.set(CMD_STOP);
    //     regs.poll_wait(SR::BUSY::SET);
    //     regs.CR_SR.set(CMD_IACK);

    //     log::info!(":::: xfer finish {error:?}, read {read}");
    //     if let Some(error) = error {
    //         Err(error)
    //     } else {
    //         Ok(read)
    //     }
    // }

    // fn i2c_write(&self, address: I2CAddress, buffer: &[u8]) -> Result<usize, Error> {
    //     let address = address.as_8_bit().ok_or(Error::NotImplemented)?;
    //     let regs = self.regs.lock();
    //     log::info!(":::: start {:#x} write", address);
    //     regs.poll_wait(SR::BUSY::SET);
    //     regs.TXR_RXR.set((address as u32) << 1);
    //     regs.CR_SR.set(CMD_START);

    //     let mut sent = 0;
    //     let mut error = None;
    //     loop {
    //         let sr = regs.poll_wait(SR::TIP::SET);
    //         if sent == buffer.len() {
    //             // Generate stop condition
    //             regs.CR_SR.set(CMD_STOP);
    //             break;
    //         }
    //         if sr.matches_all(SR::AL::SET) {
    //             regs.CR_SR.set(CMD_STOP);
    //             error = Some(Error::WouldBlock);
    //             break;
    //         }
    //         if sr.matches_all(SR::RXACK::SET) {
    //             // Generate stop condition
    //             regs.CR_SR.set(CMD_STOP);
    //             error = Some(Error::DoesNotExist);
    //             break;
    //         }

    //         log::info!(":::: send {:#x}", buffer[sent]);
    //         regs.TXR_RXR.set(buffer[sent] as u32);
    //         regs.CR_SR.set(CMD_WRITE);
    //         sent += 1;
    //     }

    //     log::info!(":::: stop {:#x}", address);
    //     regs.poll_wait(SR::BUSY::SET);
    //     regs.CR_SR.set(CMD_IACK);

    //     log::info!(":::: xfer finish {error:?}, sent {sent}");
    //     if let Some(error) = error {
    //         Err(error)
    //     } else {
    //         Ok(sent)
    //     }
    // }

    fn set_speed(&self, speed: Hertz) -> Result<Hertz, Error> {
        let input_rate = self.clock.rate()?;
        let regs = self.regs.lock();
        regs.set_clock_rate(input_rate, speed)
    }
}

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

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

            let i2c = Arc::new(I2C {
                name,
                regs: IrqSafeSpinlock::new(regs),
                clock,
                irq,
                index: OneTimeInit::new(),
                sr: BitmapEvent::new(AtomicWaker::new()),
            });

            node.make_i2c_controller(i2c.clone());

            Some(i2c)
        }
    }
}
