mod io;
mod packets;
mod types;

use io::{ReadCursor, Readable, VariableByteInteger};
use packets::*;
use pyo3::buffer::PyBuffer;
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::PyResult;
use types::*;

#[pyfunction]
fn read(py: Python, buffer: PyBuffer<u8>) -> PyResult<(Py<PyAny>, usize)> {
    // Parse the fixed header
    let mut cursor = ReadCursor::new(unsafe {
        std::slice::from_raw_parts(buffer.buf_ptr() as *const u8, buffer.len_bytes())
    });
    let first_byte = u8::read(&mut cursor)?;
    let flags = first_byte & 0x0F;
    let remaining_length = VariableByteInteger::read(&mut cursor)?;
    // Bind the cursor to the remaining bytes
    cursor.bind(remaining_length.value() as usize)?;
    // Call the read method of the corresponding packet
    #[rustfmt::skip]
    let packet = match PacketType::new(first_byte >> 4)? {
        PacketType::Connect => ConnectPacket::read(py, &mut cursor, flags)?.into(),
        PacketType::ConnAck => ConnAckPacket::read(py, &mut cursor, flags)?.into(),
        PacketType::Publish => PublishPacket::read(py, &mut cursor, flags)?.into(),
        PacketType::PubAck => PubAckPacket::read(py, &mut cursor, flags)?.into(),
        PacketType::PubRec => PubRecPacket::read(py, &mut cursor, flags)?.into(),
        PacketType::PubRel => PubRelPacket::read(py, &mut cursor, flags)?.into(),
        PacketType::PubComp => PubCompPacket::read(py, &mut cursor, flags)?.into(),
        PacketType::Subscribe => SubscribePacket::read(py, &mut cursor, flags)?.into(),
        PacketType::SubAck => SubAckPacket::read(py, &mut cursor, flags)?.into(),
        PacketType::Unsubscribe => UnsubscribePacket::read(py, &mut cursor, flags)?.into(),
        PacketType::UnsubAck => UnsubAckPacket::read(py, &mut cursor, flags)?.into(),
        PacketType::PingReq => PingReqPacket::read(py, &mut cursor, flags)?.into(),
        PacketType::PingResp => PingRespPacket::read(py, &mut cursor, flags)?.into(),
        PacketType::Disconnect => DisconnectPacket::read(py, &mut cursor, flags)?.into(),
        PacketType::Auth => AuthPacket::read(py, &mut cursor, flags)?.into(),
    };
    // Check if we've read enough bytes
    if cursor.index < cursor.buffer.len() {
        Err(PyValueError::new_err("Malformed packet"))
    } else {
        Ok((packet, cursor.index))
    }
}

#[pymodule]
fn mqtt5(m: &Bound<'_, PyModule>) -> PyResult<()> {
    // Packets
    m.add_class::<ConnectPacket>()?;
    m.add_class::<ConnAckPacket>()?;
    m.add_class::<PublishPacket>()?;
    m.add_class::<PubAckPacket>()?;
    m.add_class::<PubRecPacket>()?;
    m.add_class::<PubRelPacket>()?;
    m.add_class::<PubCompPacket>()?;
    m.add_class::<SubscribePacket>()?;
    m.add_class::<SubAckPacket>()?;
    m.add_class::<UnsubscribePacket>()?;
    m.add_class::<UnsubAckPacket>()?;
    m.add_class::<PingReqPacket>()?;
    m.add_class::<PingRespPacket>()?;
    m.add_class::<DisconnectPacket>()?;
    m.add_class::<AuthPacket>()?;
    // Reason codes
    m.add_class::<ConnAckReasonCode>()?;
    m.add_class::<PubAckReasonCode>()?;
    m.add_class::<PubRecReasonCode>()?;
    m.add_class::<PubRelReasonCode>()?;
    m.add_class::<PubCompReasonCode>()?;
    m.add_class::<SubAckReasonCode>()?;
    m.add_class::<UnsubAckReasonCode>()?;
    m.add_class::<DisconnectReasonCode>()?;
    m.add_class::<AuthReasonCode>()?;
    // Misc
    m.add_class::<QoS>()?;
    m.add_class::<RetainHandling>()?;
    m.add_class::<Will>()?;
    m.add_class::<Subscription>()?;
    // Functions
    m.add_function(wrap_pyfunction!(read, m)?)?;
    Ok(())
}
