use alloc::sync::Arc;
use device_api::{
    clock::ClockHandle,
    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},
};
use yggdrasil_abi::{
    error::Error,
    io::{TerminalOptions, TerminalOutputOptions},
};

register_bitfields! {
    u32,
    AUX_MU_IER_REG [
        RX_IRQ OFFSET(0) NUMBITS(1) [],
        TX_IRQ OFFSET(1) NUMBITS(1) [],
    ],
    AUX_MU_IIR_REG [
        IID OFFSET(1) NUMBITS(2) [
            None = 0,
            TxEmpty = 1,
            RxNotEmpty = 2,
        ],
        PENDING OFFSET(0) NUMBITS(1) [],
    ],
    AUX_MU_LSR_REG [
        TX_EMPTY OFFSET(5) NUMBITS(1) [],
    ]
}

register_structs! {
    #[allow(non_snake_case)]
    Regs {
        (0x00 => AUX_MU_IO_REG: ReadWrite<u32>),
        (0x04 => AUX_MU_IER_REG: ReadWrite<u32, AUX_MU_IER_REG::Register>),
        (0x08 => AUX_MU_IIR_REG: ReadWrite<u32, AUX_MU_IIR_REG::Register>),
        (0x0C => _0),
        (0x14 => AUX_MU_LSR_REG: ReadOnly<u32, AUX_MU_LSR_REG::Register>),
        (0x18 => _1),
        (0x30 => @END),
    }
}

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

/// Broadcom 283x mini-UART driver
pub struct Bcm2835AuxUart {
    base: PhysicalAddress,
    irq: IrqHandle,
    clock: ClockHandle,
    inner: OneTimeInit<Arc<Terminal<Inner>>>,
}

impl Regs {
    fn write_byte(&self, byte: u8) -> Result<(), Error> {
        if byte == b'\n' {
            self.write_byte(b'\r').ok();
        }

        while !self
            .AUX_MU_LSR_REG
            .matches_all(AUX_MU_LSR_REG::TX_EMPTY::SET)
        {
            core::hint::spin_loop();
        }
        self.AUX_MU_IO_REG.set(byte as u32);
        Ok(())
    }

    fn write_bytes(&self, bytes: &[u8]) -> Result<(), Error> {
        for &byte in bytes {
            self.write_byte(byte)?;
        }
        Ok(())
    }
}

impl TerminalOutput for Inner {
    fn write(&self, byte: u8) -> Result<(), Error> {
        self.regs.lock().write_byte(byte)
    }

    fn write_multiple(&self, bytes: &[u8]) -> Result<usize, Error> {
        self.regs.lock().write_bytes(bytes)?;
        Ok(bytes.len())
    }
}

impl DebugSink for Bcm2835AuxUart {
    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 InterruptHandler for Bcm2835AuxUart {
    fn handle_irq(self: Arc<Self>, _vector: IrqVector) -> bool {
        let inner = self.inner.get();

        let (status, byte) = {
            let regs = inner.output().regs.lock();

            // Reset IRQ
            regs.AUX_MU_IIR_REG.modify(AUX_MU_IIR_REG::IID::SET);

            let byte = regs.AUX_MU_IO_REG.get() as u8;
            let status = regs
                .AUX_MU_IIR_REG
                .matches_all(AUX_MU_IIR_REG::PENDING::SET);

            (status, byte)
        };

        if status {
            inner.write_to_input(byte);
        }

        status
    }
}

impl Device for Bcm2835AuxUart {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        // TODO initialize pinctrl

        self.clock.enable()?;

        let regs = unsafe { DeviceMemoryIo::map(self.base, Default::default()) }?;
        let config = TerminalOptions {
            output: TerminalOutputOptions::NL_TO_CRNL,
            ..Default::default()
        };
        let output = Inner {
            regs: IrqSafeSpinlock::new(regs),
        };
        let input = TerminalInput::with_capacity(64)?;

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

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

        Ok(())
    }

    unsafe fn init_irq(self: Arc<Self>) -> Result<(), Error> {
        self.irq.register(self.clone())?;
        self.irq.enable()?;

        let inner = self.inner.get().output();
        let regs = inner.regs.lock();

        regs.AUX_MU_IER_REG
            .modify(AUX_MU_IER_REG::RX_IRQ::SET + AUX_MU_IER_REG::TX_IRQ::CLEAR);

        Ok(())
    }

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

// TODO handle pinctrl
device_tree_driver! {
    compatible: ["brcm,bcm2835-aux-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 clock = node.clock(0)?;

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