use alloc::sync::Arc;
use device_api::{
    device::{Device, DeviceInitContext},
    interrupt::{InterruptHandler, IrqHandle, IrqVector},
};
use device_tree::driver::{device_tree_driver, Node, ProbeContext};
use libk::{
    debug::DebugSink,
    device::manager::DEVICE_REGISTRY,
    vfs::{Terminal, TerminalInput, TerminalOutput},
};
use libk_mm::{address::PhysicalAddress, device::DeviceMemoryIo};
use libk_util::{sync::IrqSafeSpinlock, OneTimeInit};
use tock_registers::{
    interfaces::{ReadWriteable, Readable, Writeable},
    register_bitfields, register_structs,
    registers::{ReadOnly, ReadWrite, WriteOnly},
};
use yggdrasil_abi::{error::Error, io::TerminalOptions};

register_bitfields! {
    u32,
    FR [
        TXFF OFFSET(5) NUMBITS(1) [],
        RXFE OFFSET(4) NUMBITS(1) [],
        BUSY OFFSET(3) NUMBITS(1) [],
    ],
    CR [
        RXE OFFSET(9) NUMBITS(1) [],
        TXE OFFSET(8) NUMBITS(1) [],
        UARTEN OFFSET(0) NUMBITS(1) [],
    ],
    ICR [
        ALL OFFSET(0) NUMBITS(11) [],
    ],
    IMSC [
        RXIM OFFSET(4) NUMBITS(1) [],
    ]
}

register_structs! {
    #[allow(non_snake_case)]
    Regs {
        /// Transmit/receive data register
        (0x00 => DR: ReadWrite<u32>),
        (0x04 => _0),
        (0x18 => FR: ReadOnly<u32, FR::Register>),
        (0x1C => _1),
        (0x2C => LCR_H: ReadWrite<u32>),
        (0x30 => CR: ReadWrite<u32, CR::Register>),
        (0x34 => IFLS: ReadWrite<u32>),
        (0x38 => IMSC: ReadWrite<u32, IMSC::Register>),
        (0x3C => _2),
        (0x44 => ICR: WriteOnly<u32, ICR::Register>),
        (0x48 => @END),
    }
}

struct Io {
    regs: DeviceMemoryIo<'static, Regs>,
}

struct Pl011Inner {
    io: IrqSafeSpinlock<Io>,
}

/// PL011 device instance
pub struct Pl011 {
    inner: OneTimeInit<Arc<Terminal<Pl011Inner>>>,
    base: PhysicalAddress,
    irq: IrqHandle,
}

impl Io {
    fn send(&mut self, byte: u8) {
        while self.regs.FR.matches_all(FR::TXFF::SET) {
            core::hint::spin_loop();
        }
        self.regs.DR.set(byte as u32);
    }

    unsafe fn init(&mut self) {
        self.regs.CR.set(0);
        self.regs.ICR.write(ICR::ALL::CLEAR);
        self.regs
            .CR
            .write(CR::UARTEN::SET + CR::TXE::SET + CR::RXE::SET);
    }
}

impl TerminalOutput for Pl011Inner {
    fn write(&self, byte: u8) -> Result<(), Error> {
        self.io.lock().send(byte);
        Ok(())
    }

    fn write_multiple(&self, bytes: &[u8]) -> Result<usize, Error> {
        let mut lock = self.io.lock();
        for &byte in bytes {
            lock.send(byte);
        }
        Ok(bytes.len())
    }
}

impl Pl011Inner {
    #[inline]
    fn handle_irq(&self) -> u32 {
        let lock = self.io.lock();
        lock.regs.ICR.write(ICR::ALL::CLEAR);
        lock.regs.DR.get()
    }
}

impl DebugSink for Pl011 {
    fn putc(&self, byte: u8) -> Result<(), Error> {
        self.inner.get().putc_to_output(byte)
    }

    fn supports_control_sequences(&self) -> bool {
        true
    }
}

impl InterruptHandler for Pl011 {
    fn handle_irq(self: Arc<Self>, _vector: IrqVector) -> bool {
        let inner = self.inner.get();
        let byte = inner.output().handle_irq();

        inner.write_to_input(byte as u8);

        true
    }
}

impl Device for Pl011 {
    fn display_name(&self) -> &'static str {
        "Primecell PL011 UART"
    }

    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        let mut io = Io {
            regs: DeviceMemoryIo::map(self.base, Default::default())?,
        };
        io.init();

        let input = TerminalInput::with_capacity(64)?;
        let output = Pl011Inner {
            io: IrqSafeSpinlock::new(io),
        };

        let terminal = self.inner.init(Arc::new(Terminal::from_parts(
            TerminalOptions::const_default(),
            input,
            output,
        )));

        DEVICE_REGISTRY
            .serial_terminal
            .register(terminal.clone(), Some(self.clone()))
            .ok();

        Ok(())
    }

    unsafe fn init_irq(self: Arc<Self>) -> Result<(), Error> {
        self.irq.register(self.clone())?;
        {
            let io = self.inner.get().output().io.lock();
            io.regs.IMSC.modify(IMSC::RXIM::SET);
        }
        self.irq.enable()?;

        Ok(())
    }
}

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

            Some(Arc::new(Pl011 {
                base,
                irq,
                inner: OneTimeInit::new(),
            }))
        }
    }
}
