//! 16550-style UART device driver
use abi::{error::Error, io::TerminalOptions};
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},
};

register_bitfields!(
    u8,
    IER [
        /// Received data ready
        RDR OFFSET(0) NUMBITS(1) [],
        /// Trasmitter holding register empty
        THRE OFFSET(1) NUMBITS(1) [],
        /// Receiver line status
        RLS OFFSET(2) NUMBITS(1) [],
        /// Modem status
        MS OFFSET(3) NUMBITS(1) [],
    ],
    LSR [
        /// Data ready indicator
        DR OFFSET(0) NUMBITS(1) [],
        /// Transmitter FIFO empty
        TFE OFFSET(5) NUMBITS(1) [],
    ],
    LCR [
        BITS OFFSET(0) NUMBITS(2) [
            Bits8 = 3
        ],
        STOPBITS OFFSET(2) NUMBITS(1) [],
        PARITY OFFSET(3) NUMBITS(1) [],
        PARITY_EVEN OFFSET(4) NUMBITS(1) [],
        PARITY_STICK OFFSET(5) NUMBITS(1) [],
        BREAK OFFSET(6) NUMBITS(1) [],
        DLAB OFFSET(7) NUMBITS(1) [],
    ]
);

register_structs! {
    #[allow(non_snake_case)]
    Regs {
        // Read: receive buffer, write: transmit buffer
        (0x00 => DR: ReadWrite<u8>),
        (0x01 => IER: ReadWrite<u8, IER::Register>),
        // Read: interrupt idenditication, write: FIFO control
        (0x02 => FCR: ReadWrite<u8>),
        (0x03 => LCR: ReadWrite<u8, LCR::Register>),
        (0x04 => MCR: WriteOnly<u8>),
        (0x05 => LSR: ReadOnly<u8, LSR::Register>),
        (0x06 => MSR: ReadOnly<u8>),
        (0x07 => _0),
        (0x08 => @END),
    }
}

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

struct Inner {
    io: IrqSafeSpinlock<Io>,
}

/// ns16550a-style UART driver
pub struct Ns16550a {
    inner: OneTimeInit<Arc<Terminal<Inner>>>,
    base: PhysicalAddress,
    irq: IrqHandle,
}

impl Io {
    fn init(&mut self) {
        self.regs.LCR.write(
            LCR::BITS::Bits8 + LCR::BREAK::CLEAR + LCR::STOPBITS::CLEAR + LCR::PARITY::CLEAR,
        );
        self.regs.IER.set(0);
    }

    fn send(&mut self, byte: u8) {
        while self.regs.LSR.matches_all(LSR::TFE::CLEAR) {
            core::hint::spin_loop();
        }
        self.regs.DR.set(byte);
    }

    fn handle_irq(&self) -> Option<u8> {
        let status = self.regs.FCR.get() & 0xF;
        if status == 0b1100 || status == 0b0100 {
            Some(self.regs.DR.get())
        } else {
            None
        }
    }
}

impl Device for Ns16550a {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        log::info!("Init ns16550a @ {:#x}", self.base);
        let mut io = Io {
            regs: DeviceMemoryIo::map(self.base, Default::default())?,
        };
        io.init();

        let input = TerminalInput::with_capacity(64)?;
        let output = Inner {
            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())?;
        self.irq.enable()?;
        let io = self.inner.get().output().io.lock();
        io.regs.IER.modify(IER::RDR::SET);
        Ok(())
    }

    fn display_name(&self) -> &str {
        "ns16550a UART"
    }
}

impl InterruptHandler for Ns16550a {
    fn handle_irq(self: Arc<Self>, _vector: IrqVector) -> bool {
        let inner = self.inner.get();
        let output = inner.output();
        let byte = output.io.lock().handle_irq();
        if let Some(byte) = byte {
            inner.write_to_input(byte);
            true
        } else {
            false
        }
    }
}

impl TerminalOutput for Inner {
    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 DebugSink for Ns16550a {
    fn putc(&self, c: u8) -> Result<(), Error> {
        self.inner.get().putc_to_output(c)
    }

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

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

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