//! Synopsys DesignWare 8250 driver
use abi::{error::Error, io::TerminalOptions};
use alloc::sync::Arc;
use device_api::{
    clock::{ClockHandle, Hertz, ResetHandle},
    device::{Device, DeviceInitContext},
    interrupt::{InterruptHandler, IrqHandle, IrqVector},
};
use device_tree::driver::{device_tree_driver, Node, ProbeContext};
use libk::{
    debug::{self, 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! {
    u32,
    LCR [
        DLAB    OFFSET(7)   NUMBITS(1)  [],
        DLS     OFFSET(0)   NUMBITS(2)  [
            DL8b = 0b11,
        ],
    ],
    IER [
        PTIME   OFFSET(7)   NUMBITS(1)  [],
        EDSSI   OFFSET(3)   NUMBITS(1)  [],
        ELSI    OFFSET(2)   NUMBITS(1)  [],
        // Transmit buffer available
        ETBEI   OFFSET(1)   NUMBITS(1)  [],
        // Receive data available
        ERBFI   OFFSET(0)   NUMBITS(1)  [],
    ],
    LSR [
        // Data ready bit
        DR      OFFSET(0)   NUMBITS(1)  [],
        // Transmitter holding register empty
        THRE    OFFSET(5)   NUMBITS(1)  [],
    ],
    USR [
        BUSY    OFFSET(0)   NUMBITS(1)  [],
    ]
}

register_structs! {
    #[allow(non_snake_case)]
    Regs {
        // DLAB=0, Write: transmitter holding register/Read: receiver buffer register
        // DLAB=1, Read/Write: divisor latch low
        (0x000 => DR: ReadWrite<u32>),
        // DLAB=0: Interrupt enable register
        // DLAB=1: Divisor latch high
        (0x004 => IER: ReadWrite<u32, IER::Register>),
        // Read: interrupt identification register/Write: frame control register
        (0x008 => IIR: ReadWrite<u32>),
        // Line control register
        (0x00C => LCR: ReadWrite<u32, LCR::Register>),
        // Modem control register
        (0x010 => MCR: ReadWrite<u32>),
        // Line status register
        (0x014 => LSR: ReadOnly<u32, LSR::Register>),
        // Modem status register
        (0x018 => MSR: ReadOnly<u32>),
        // Scratchpad
        (0x01C => SCR: ReadWrite<u32>),
        // Low-power divisor latch low
        (0x020 => LPDLL: ReadWrite<u32>),
        // Low-power divisor latch high
        (0x024 => LPDLH: ReadWrite<u32>),
        (0x028 => _0),
        // Shadow receive/transmit buffer
        (0x030 => SDR: [ReadWrite<u32>; 16]),
        (0x070 => FAR: ReadWrite<u32>),
        (0x074 => TFR: ReadOnly<u32>),
        (0x078 => RFW: WriteOnly<u32>),
        (0x07C => USR: ReadOnly<u32, USR::Register>),
        (0x080 => TFL: ReadOnly<u32>),
        (0x084 => RFL: ReadOnly<u32>),
        (0x088 => SRR: WriteOnly<u32>),
        (0x08C => SRTS: ReadWrite<u32>),
        (0x090 => SBCR: ReadWrite<u32>),
        (0x094 => SDMAM: ReadWrite<u32>),
        (0x098 => SFE: ReadWrite<u32>),
        (0x09C => SRT: ReadWrite<u32>),
        (0x0A0 => STET: ReadWrite<u32>),
        (0x0A4 => HTX: ReadWrite<u32>),
        (0x0A8 => DMASA: WriteOnly<u32>),
        (0x0AC => _1),
        (0x0F4 => CPR: ReadOnly<u32>),
        (0x0F8 => UCV: ReadOnly<u32>),
        (0x0FC => CTR: ReadOnly<u32>),
        (0x100 => @END),
    }
}

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

struct Inner {
    io: IrqSafeSpinlock<Io>,
}

/// Synopsys DesignWare 8250 UART
pub struct DwUart {
    base: PhysicalAddress,
    irq: IrqHandle,
    clk_baud: ClockHandle,
    #[allow(unused)]
    clk_apb: Option<ClockHandle>,
    rst: Option<ResetHandle>,

    inner: OneTimeInit<Arc<Terminal<Inner>>>,
}

impl Io {
    fn send(&mut self, byte: u8) {
        // TODO
        if byte == b'\n' {
            self.send(b'\r');
        }

        while !self.regs.LSR.matches_all(LSR::THRE::SET) {
            core::hint::spin_loop();
        }
        self.regs.DR.set(byte as u32);
    }

    fn init(&mut self, baud_clock: Hertz, baud_rate: u64) {
        let divisor = (baud_clock.0 / (baud_rate * 16)) as u32;

        self.wait_busy();
        self.regs.IER.set(0);
        for _ in 0..100 {
            let _ = self.regs.LSR.get();
        }
        for _ in 0..100 {
            let _ = self.regs.DR.get();
        }
        self.wait_busy();
        self.regs.LCR.write(LCR::DLAB::SET);
        self.wait_busy();
        self.regs.DR.set(divisor & 0xFF);
        self.regs.IER.set((divisor >> 8) & 0xFF);
        self.wait_busy();
        self.regs.LCR.write(LCR::DLS::DL8b);
        self.wait_busy();
        self.regs.IIR.set(0x01);
        self.wait_busy();
        self.regs.MCR.set(0x00);
        let _ = self.regs.LSR.get();
        for _ in 0..100 {
            let _ = self.regs.DR.get();
        }
        self.regs.SCR.set(0x00);
    }

    fn handle_irq(&mut self) -> Option<u8> {
        let status = self.regs.IIR.get();

        if status & 0xF == 4 {
            Some(self.regs.DR.get() as u8)
        } else {
            None
        }
    }

    fn wait_busy(&self) {
        for _ in 0..100000 {
            core::hint::spin_loop();
        }
        let mut timeout = 1000000;
        while timeout > 0 && self.regs.USR.matches_all(USR::BUSY::SET) {
            core::hint::spin_loop();
            timeout -= 1;
        }
    }
}

impl InterruptHandler for DwUart {
    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 Device for DwUart {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        let baud_rate = 115200;
        let clk_baud_rate = self.clk_baud.rate()?;

        // Prevent firmware (SBI in riscv64) from printing to UART while it's being
        // reset/initialized
        let guard = debug::MuteGuard::acquire();

        let regs = DeviceMemoryIo::<Regs>::map(self.base, Default::default())?;
        let mut io = Io { regs };

        if let Some(reset) = self.rst.as_ref() {
            reset.assert_for_cycles(100000)?;
        }

        io.init(clk_baud_rate, baud_rate);

        io.send(b'\r');
        io.send(b'\n');

        drop(guard);

        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 output = self.inner.get().output();
        let io = output.io.lock();
        io.regs.IER.modify(IER::ERBFI::SET);
        Ok(())
    }

    fn display_name(&self) -> &str {
        "Synopsys DesignWare 8250 UART"
    }
}

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

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

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())
    }
}

device_tree_driver! {
    compatible: ["snps,dw-apb-uart"],
    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)?;
            let clk_baud = node.named_clock("baudclk")?;
            let clk_apb = node.named_clock("apb_pclk");
            let rst = node.reset(0);

            Some(Arc::new(DwUart {
                base,
                irq,
                clk_baud,
                clk_apb,
                rst,

                inner: OneTimeInit::new()
            }))
        }
    }
}
