//! Driver for x86 COM ports
use abi::{error::Error, io::TerminalOptions};
use alloc::sync::Arc;
use device_api::{
    device::{Device, DeviceInitContext},
    interrupt::{InterruptHandler, Irq, IrqVector},
};
use kernel_arch_x86::intrinsics::{IoPort, IoPortAccess};
use libk::{
    debug::DebugSink,
    device::{external_interrupt_controller, manager::DEVICE_REGISTRY},
    vfs::{Terminal, TerminalInput, TerminalOutput},
};
use libk_util::sync::IrqSafeSpinlock;

// Single port
struct Regs {
    dr: IoPort<u8>,
    lsr: IoPort<u8>,
    lcr: IoPort<u8>,
    ier: IoPort<u8>,
    isr: IoPort<u8>,

    baud_rate: u32,
}

struct PortInner {
    regs: IrqSafeSpinlock<Regs>,
}

/// Single port of the COM port pair
pub struct Port {
    terminal: Arc<Terminal<PortInner>>,
    irq: Irq,
}

/// COM port pair
#[allow(unused)]
pub struct ComPort {
    port_a: Arc<Port>,
    port_b: Arc<Port>,
}

impl Regs {
    fn write(&mut self, byte: u8) -> Result<(), Error> {
        while self.lsr.read() & Port::LSR_THRE == 0 {
            core::hint::spin_loop();
        }

        self.dr.write(byte);
        Ok(())
    }
}

impl PortInner {
    fn handle_irq(&self) -> Option<u8> {
        let (status, value) = {
            let inner = self.regs.lock();
            let status = inner.isr.read();
            let value = inner.dr.read();
            (status, value)
        };

        if status & Port::ISR_IRQ_MASK != 0 {
            Some(value)
        } else {
            None
        }
    }
}

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

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

    fn baud_rate(&self) -> u32 {
        self.regs.lock().baud_rate
    }

    fn set_baud_rate(&self, baud: u32) -> Result<(), Error> {
        if baud > 115200 {
            log::warn!("Tried to set baud rate for COM port beyond 115200: {baud}");
            return Err(Error::InvalidArgument);
        }

        let div = 115200 / baud;

        let mut regs = self.regs.lock();

        regs.lcr.write(regs.lcr.read() | (1 << 7));
        regs.dr.write(div as u8);
        regs.ier.write((div >> 8) as u8);
        regs.lcr.write(regs.lcr.read() & !(1 << 7));

        regs.baud_rate = baud;

        Ok(())
    }
}

impl DebugSink for Port {
    fn putc(&self, c: u8) -> Result<(), Error> {
        self.terminal.putc_to_output(c).ok();
        Ok(())
    }

    fn puts(&self, s: &str) -> Result<(), Error> {
        self.terminal.write_to_output(s.as_bytes()).ok();
        Ok(())
    }

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

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

impl Device for Port {
    fn display_name(&self) -> &'static str {
        "COM port"
    }

    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        self.terminal.output().set_baud_rate(115200)?;

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

    unsafe fn init_irq(self: Arc<Self>) -> Result<(), Error> {
        let intc = external_interrupt_controller()?;

        // TODO check that the same IRQ is not bound already for another port
        intc.register_irq(self.irq, Default::default(), self.clone())?;
        intc.enable_irq(self.irq)?;

        let regs = self.terminal.output().regs.lock();

        regs.ier.modify(|v| v | Self::IER_RXDA);

        Ok(())
    }
}

impl Port {
    const LSR_THRE: u8 = 1 << 5;
    const IER_RXDA: u8 = 1 << 0;
    const ISR_IRQ_MASK: u8 = 3 << 1;

    fn new(base: u16, irq: Irq) -> Result<Self, Error> {
        let input = TerminalInput::with_capacity(64)?;
        let output = PortInner {
            regs: IrqSafeSpinlock::new(Regs {
                dr: IoPort::new(base),
                lsr: IoPort::new(base + 5),
                ier: IoPort::new(base + 1),
                isr: IoPort::new(base + 2),
                lcr: IoPort::new(base + 3),

                baud_rate: 0,
            }),
        };
        let terminal = Terminal::from_parts(TerminalOptions::const_default(), input, output);
        Ok(Self {
            terminal: Arc::new(terminal),
            irq,
        })
    }
}

impl ComPort {
    /// Constructs a COM port pair
    fn new(port_a: u16, port_b: u16, irq: Irq) -> Result<Self, Error> {
        Ok(Self {
            port_a: Arc::new(Port::new(port_a, irq)?),
            port_b: Arc::new(Port::new(port_b, irq)?),
        })
    }

    pub fn setup(
        cx: DeviceInitContext,
        port_a: u16,
        port_b: u16,
        irq: Irq,
    ) -> Result<Arc<Self>, Error> {
        let this = Arc::new(Self::new(port_a, port_b, irq)?);
        unsafe { this.port_a().clone().init(cx) }?;
        Ok(this)
    }

    /// Returns a reference to the A port of this COM pair
    pub fn port_a(&self) -> &Arc<Port> {
        &self.port_a
    }
}
