#![no_std]

#[derive(Debug, PartialEq, Eq)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
pub enum SntpError {
    InvalidPacket,
    InvalidVersion,
    InvalidMode,
    KissOfDeath,
    InvalidTime,
    NetFailure,
    NetTimeout,
}

pub struct SntpRequest;

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
#[cfg_attr(feature = "defmt", derive(defmt::Format))]
/// SNTP timestamp.
pub struct SntpTimestamp(u64);

impl SntpTimestamp {
    pub fn microseconds(&self) -> u64 {
        (self.0 >> 32) * 1_000_000 + ((self.0 & 0xFFFFFFFF) * 1_000_000 / 0x100000000)
    }

    /// Returns true if the most significant bit is set.
    ///
    /// Relevant documentation from RFC 2030:
    ///
    /// ```text
    /// Note that, since some time in 1968 (second 2,147,483,648) the most
    /// significant bit (bit 0 of the integer part) has been set and that the
    /// 64-bit field will overflow some time in 2036 (second 4,294,967,296).
    /// Should NTP or SNTP be in use in 2036, some external means will be
    /// necessary to qualify time relative to 1900 and time relative to 2036
    /// (and other multiples of 136 years). There will exist a 200-picosecond
    /// interval, henceforth ignored, every 136 years when the 64-bit field
    /// will be 0, which by convention is interpreted as an invalid or
    /// unavailable timestamp.
    ///    As the NTP timestamp format has been in use for the last 17 years,
    ///    it remains a possibility that it will be in use 40 years from now
    ///    when the seconds field overflows. As it is probably inappropriate
    ///    to archive NTP timestamps before bit 0 was set in 1968, a
    ///    convenient way to extend the useful life of NTP timestamps is the
    ///    following convention: If bit 0 is set, the UTC time is in the
    ///    range 1968-2036 and UTC time is reckoned from 0h 0m 0s UTC on 1
    ///    January 1900. If bit 0 is not set, the time is in the range 2036-
    ///    2104 and UTC time is reckoned from 6h 28m 16s UTC on 7 February
    ///    2036. Note that when calculating the correspondence, 2000 is not a
    ///    leap year. Note also that leap seconds are not counted in the
    ///    reckoning.
    ///```
    pub fn msb_set(&self) -> bool {
        self.0 & (1 << 63) != 0
    }

    /// Microseconds since the UNIX epoch.
    pub fn utc_micros(&self) -> i64 {
        let ntp_epoch_micros = self.microseconds() as i64;
        let offset: i64 = if self.msb_set() {
            -2208988800000000
        } else {
            2085978496000000
        };

        ntp_epoch_micros + offset
    }

    #[cfg(feature = "chrono")]
    pub fn try_to_naive_datetime(self) -> Result<chrono::NaiveDateTime, SntpError> {
        self.try_into()
    }
}

impl SntpRequest {
    pub const SNTP_PACKET_SIZE: usize = 48;

    pub const fn create_packet() -> [u8; Self::SNTP_PACKET_SIZE] {
        let mut packet = [0u8; Self::SNTP_PACKET_SIZE];
        packet[0] = (3 << 6) | (4 << 3) | 3;
        packet
    }

    pub fn create_packet_from_buffer(packet: &mut [u8]) -> Result<(), SntpError> {
        if packet.len() != Self::SNTP_PACKET_SIZE {
            return Err(SntpError::InvalidPacket);
        }

        packet[0] = (3 << 6) | (4 << 3) | 3;

        Ok(())
    }

    pub fn read_timestamp(packet: &[u8]) -> Result<SntpTimestamp, SntpError> {
        if packet.len() == Self::SNTP_PACKET_SIZE {
            let header = packet[0];
            let version = (header & 0x38) >> 3;

            if version != 4 {
                return Err(SntpError::InvalidVersion);
            }

            let mode = header & 0x7;

            if !(4..=5).contains(&mode) {
                return Err(SntpError::InvalidMode);
            }

            let kiss_of_death = packet[1] == 0;

            if kiss_of_death {
                return Err(SntpError::KissOfDeath);
            }

            let timestamp = SntpTimestamp(read_be_u64(&packet[40..48]));

            return Ok(timestamp);
        }

        Err(SntpError::InvalidPacket)
    }

    #[cfg(feature = "chrono")]
    pub fn as_naive_datetime(raw_time: SntpTimestamp) -> Result<chrono::NaiveDateTime, SntpError> {
        raw_time.try_into()
    }
}

#[cfg(feature = "chrono")]
impl TryFrom<SntpTimestamp> for chrono::NaiveDateTime {
    type Error = SntpError;

    fn try_from(timestamp: SntpTimestamp) -> Result<Self, Self::Error> {
        Ok(chrono::DateTime::<chrono::Utc>::try_from(timestamp)?.naive_utc())
    }
}

#[cfg(feature = "chrono")]
impl TryFrom<SntpTimestamp> for chrono::DateTime<chrono::Utc> {
    type Error = SntpError;

    fn try_from(timestamp: SntpTimestamp) -> Result<Self, Self::Error> {
        chrono::DateTime::<chrono::Utc>::from_timestamp_micros(timestamp.utc_micros())
            .ok_or(SntpError::InvalidTime)
    }
}

#[inline]
fn read_be_u64(input: &[u8]) -> u64 {
    let (int_bytes, _) = input.split_at(core::mem::size_of::<u64>());
    u64::from_be_bytes(int_bytes.try_into().unwrap())
}

#[cfg(feature = "embassy-net")]
pub trait SntpSocket {
    fn resolve_time(
        &mut self,
        addrs: &[embassy_net::IpAddress],
    ) -> impl Future<Output = Result<SntpTimestamp, SntpError>>;
}

#[cfg(feature = "embassy-net")]
impl SntpSocket for embassy_net::udp::UdpSocket<'_> {
    async fn resolve_time(
        &mut self,
        addrs: &[embassy_net::IpAddress],
    ) -> Result<SntpTimestamp, SntpError> {
        use embassy_time::{Duration, WithTimeout};
        use sachy_fmt::{debug, error};

        if addrs.is_empty() {
            return Err(SntpError::NetFailure);
        }

        debug!("SNTP address list: {}", addrs);

        let addr = addrs[0];

        debug!("Binding to port 123");
        self.bind(123).map_err(|_| SntpError::NetFailure)?;

        debug!("Sending SNTP request");
        self.send_to_with(
            SntpRequest::SNTP_PACKET_SIZE,
            embassy_net::IpEndpoint::new(addr, 123),
            SntpRequest::create_packet_from_buffer,
        )
        .await
        .map_err(|e| {
            error!("Failed to send: {}", e);
            SntpError::NetFailure
        })??;

        debug!("Waiting for a response...");
        let res = self
            .recv_from_with(|buf, _from| SntpRequest::read_timestamp(buf))
            .with_timeout(Duration::from_secs(5))
            .await
            .map_err(|_| SntpError::NetTimeout)
            .flatten();

        debug!("Received: {}", &res);
        debug!("Closing SNTP port...");
        self.close();

        res
    }
}