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

register_bitfields! {
    u32,
    txctrl [
        txen    OFFSET(0)   NUMBITS(1)  [],
        nstop   OFFSET(1)   NUMBITS(1)  [
            One = 0,
            Two = 1
        ],
        txcnt   OFFSET(16)  NUMBITS(3)  [],
    ],
    rxctrl [
        rxen    OFFSET(0)   NUMBITS(1)  [],
        rxcnt   OFFSET(16)  NUMBITS(3)  [],
    ],
    interrupt [
        txwm    OFFSET(0)   NUMBITS(1)  [],
        rxwm    OFFSET(1)   NUMBITS(1)  [],
    ],
}

register_structs! {
    Regs {
        (0x00 => txdata: WriteOnly<u32>),
        (0x04 => rxdata: ReadOnly<u32>),
        (0x08 => txctrl: ReadWrite<u32, txctrl::Register>),
        (0x0C => rxctrl: ReadWrite<u32, rxctrl::Register>),
        (0x10 => ie: ReadWrite<u32, interrupt::Register>),
        (0x14 => ip: ReadOnly<u32, interrupt::Register>),
        (0x18 => div: ReadWrite<u32>),
        (0x1C => _0),
        (0x20 => @END),
    }
}

struct Inner {
    clock: ClockHandle,
    regs: IrqSafeSpinlock<DeviceMemoryIo<'static, Regs>>,
}

struct Uart {
    name: &'static str,
    irq: IrqHandle,
    clock: ClockHandle,
    base: PhysicalAddress,
    inner: OneTimeInit<Arc<Terminal<Inner>>>,
}

impl Regs {
    fn set_baud_rate(&self, input_clk: Hertz, baud: u32) -> Result<(), Error> {
        let div = Hertz::divider(input_clk, Hertz::from(baud)).ok_or(Error::InvalidArgument)?;
        if div >= (1 << 20) || div == 0 {
            return Err(Error::InvalidArgument);
        }
        let div = div - 1;
        self.div.set(div);

        Ok(())
    }

    fn baud_rate(&self, input_clk: Hertz) -> u32 {
        let div = self.div.get() + 1;
        (input_clk / div).0 as u32
    }

    fn write(&self, byte: u8) {
        while !self.ip.matches_all(interrupt::txwm::SET) {
            core::hint::spin_loop();
        }
        self.txdata.set(byte as u32);
    }
}

impl Device for Uart {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        let regs = unsafe { DeviceMemoryIo::<Regs>::map(self.base, Default::default()) }?;

        self.clock.enable()?;
        let input_clk = self.clock.rate()?;

        {
            let _guard = debug::MuteGuard::acquire();

            regs.txctrl.write(txctrl::txen::CLEAR);
            regs.rxctrl.write(rxctrl::rxen::CLEAR);
            regs.set_baud_rate(input_clk, 115200)?;
            regs.txctrl.write(txctrl::txen::SET + txctrl::txcnt.val(3));
            regs.rxctrl.write(rxctrl::rxen::SET + rxctrl::rxcnt.val(0));
        }

        let input = TerminalInput::with_capacity(64)?;
        let output = Inner {
            regs: IrqSafeSpinlock::new(regs),
            clock: self.clock.clone(),
        };

        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 regs = self.inner.get().output().regs.lock();
        regs.ie.write(interrupt::rxwm::SET);
        Ok(())
    }

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

impl InterruptHandler for Uart {
    fn handle_irq(self: Arc<Self>, _vector: IrqVector) -> bool {
        let terminal = self.inner.get();
        let byte = {
            let regs = terminal.output().regs.lock();
            if regs.ip.matches_all(interrupt::rxwm::SET) {
                regs.rxdata.get() as u8
            } else {
                return false;
            }
        };
        terminal.write_to_input(byte);
        true
    }
}

impl DebugSink for Uart {
    fn putc(&self, c: u8) -> Result<(), Error> {
        self.inner.get().putc_to_output(c)
    }
    fn puts(&self, s: &str) -> Result<(), Error> {
        self.inner.get().write_to_output(s.as_bytes())?;
        Ok(())
    }
    fn supports_control_sequences(&self) -> bool {
        true
    }
}

impl TerminalOutput for Inner {
    fn write(&self, byte: u8, options: &TerminalOutputOptions) -> Result<(), Error> {
        let regs = self.regs.lock();
        if byte == b'\n' && options.contains(TerminalOutputOptions::NL_TO_CRNL) {
            regs.write(b'\r');
        }
        regs.write(byte);
        Ok(())
    }
    fn write_multiple(
        &self,
        bytes: &[u8],
        options: &TerminalOutputOptions,
    ) -> Result<usize, Error> {
        let regs = self.regs.lock();
        for &byte in bytes {
            if byte == b'\n' && options.contains(TerminalOutputOptions::NL_TO_CRNL) {
                regs.write(b'\r');
            }
            regs.write(byte);
        }
        Ok(bytes.len())
    }

    fn baud_rate(&self) -> u32 {
        let input_clk = self.clock.rate().unwrap();
        self.regs.lock().baud_rate(input_clk)
    }
    fn set_baud_rate(&self, baud: u32) -> Result<(), Error> {
        let input_clk = self.clock.rate().unwrap();
        self.regs.lock().set_baud_rate(input_clk, baud)
    }
}

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

            let uart = Arc::new(Uart {
                name,
                irq,
                clock,
                base,
                inner: OneTimeInit::new()
            });

            Some(uart)
        }
    }
}
