// Copyright (c) 2020 Xu Shaohua <shaohua@biofan.org>. All rights reserved.
// Use of this source is governed by General Public License that can be found
// in the LICENSE file.

#![no_std]

extern crate alloc;
extern crate core;
use alloc::vec::Vec;
use byteorder::{BigEndian, WriteBytesExt};
use core::convert::TryInto;
use core::default::Default;

/// Number of bytes of binary representation of struct.
pub trait BinLen {
    const LEN: usize;
}

#[derive(Debug)]
pub enum Error {
    InvalidPacket,
}

pub trait ToNetPacket {
    fn to_net(&self, v: &mut Vec<u8>);
}

/// 2-bits indicator to check leap second.
#[repr(u8)]
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum LeapIndicator {
    NoWarning = 0,
    AddOne = 1,
    SubOne = 2,
    Unknown = 3,
}

impl Default for LeapIndicator {
    fn default() -> Self {
        LeapIndicator::NoWarning
    }
}

/// 3-bits version number of NTP protocol.
#[repr(u8)]
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Version {
    V1 = 0,
    V2 = 1,
    V3 = 3,
    V4 = 4,
    V5 = 5,
}

impl Default for Version {
    fn default() -> Self {
        Version::V4
    }
}

/// 3-bits mode type of NTP packet.
#[repr(u8)]
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Mode {
    Reserved = 0,
    SymmetricActive = 1,
    SymmetricPassive = 2,
    Client = 3,
    Server = 4,
    Broadcast = 5,
    ControlMessage = 6,
    ReservedForPrivateUse = 7,
}

impl Default for Mode {
    fn default() -> Self {
        Mode::Client
    }
}

/// 8-bit stratum type.
#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
pub struct Stratum(u8);

/// NTP packet header.
#[repr(C)]
#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
pub struct Header {
    /// Flags(the first byte) are constructed with leap_id, version and mode.
    pub leap_id: LeapIndicator,
    pub version: Version,
    pub mode: Mode,

    pub stratum: Stratum,
    pub poll: i8,
    pub precision: i8,
}

impl BinLen for Header {
    const LEN: usize = 4;
}

impl ToNetPacket for Header {
    fn to_net(&self, v: &mut Vec<u8>) {
        let flags = ((self.leap_id as u8) & 0b11000000)
            + ((self.version as u8) << 3)
            + ((self.mode as u8) & 0b00000111);
        log::info!("flags: {:x}", flags);
        v.push(flags);
        v.push(self.stratum.0);
        // TODO(Shaohua): Check interger sign
        v.push(self.poll as u8);
        v.push(self.precision as u8);
    }
}

/// Short format of timestamp.
#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
pub struct ShortTimestamp {
    pub seconds: u16,
    pub fraction: u16,
}

impl BinLen for ShortTimestamp {
    const LEN: usize = 4;
}

impl ToNetPacket for ShortTimestamp {
    fn to_net(&self, v: &mut Vec<u8>) {
        v.write_u16::<BigEndian>(self.seconds);
        v.write_u16::<BigEndian>(self.fraction);
    }
}

/// Reference identifier.
#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
pub struct RefIdentifier {
    pub ref_id: [u8; 4],
}

impl BinLen for RefIdentifier {
    const LEN: usize = 4;
}

impl ToNetPacket for RefIdentifier {
    fn to_net(&self, v: &mut Vec<u8>) {
        v.push(self.ref_id[0]);
        v.push(self.ref_id[1]);
        v.push(self.ref_id[2]);
        v.push(self.ref_id[3]);
    }
}

/// 64-bits timestamp.
#[repr(C)]
#[derive(Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct Timestamp {
    pub seconds: u32,
    pub fraction: u32,
}

impl BinLen for Timestamp {
    const LEN: usize = 8;
}

impl ToNetPacket for Timestamp {
    fn to_net(&self, v: &mut Vec<u8>) {
        v.write_u32::<BigEndian>(self.seconds);
        v.write_u32::<BigEndian>(self.fraction);
    }
}

impl Timestamp {
    pub fn now() -> Timestamp {
        let mut now = nc::timeval_t::default();
        let mut tz = nc::timezone_t::default();
        let _ = nc::gettimeofday(&mut now, &mut tz);
        const DELTA: u64 = 2_208_988_800;
        let seconds = (now.tv_sec as u64 + DELTA).try_into().unwrap();
        let fraction = (now.tv_usec as u64).try_into().unwrap();
        Timestamp { seconds, fraction }
    }
}

/// ## Ntp package format:
/// ```memory
///  0                   1                   2                   3
///  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
/// +--------------+---------------+---------------+----------------+
/// |LI | VN  |Mode |    Stratum     |     Poll      |  Precision   |
/// +--------------+---------------+---------------+----------------+
/// |                          Root Delay                           |
/// +--------------+---------------+---------------+----------------+
/// |                        Root Dispersion                        |
/// +--------------+---------------+---------------+----------------+
/// |                         Reference ID                          |
/// +--------------+---------------+---------------+----------------+
/// |                                                               |
/// +                   Reference Timestamp (64)                    +
/// |                                                               |
/// +--------------+---------------+---------------+----------------+
/// |                                                               |
/// +                     Origin Timestamp (64)                     +
/// |                                                               |
/// +--------------+---------------+---------------+----------------+
/// |                                                               |
/// +                    Receive Timestamp (64)                     +
/// |                                                               |
/// +--------------+---------------+---------------+----------------+
/// |                                                               |
/// +                    Transmit Timestamp (64)                    +
/// |                                                               |
/// +--------------+---------------+---------------+----------------+
/// |                                                               |
/// .                                                               .
/// .                    Extension Field 1 (variable)               .
/// .                                                               .
/// |                                                               |
/// +--------------+---------------+---------------+----------------+
/// |                                                               |
/// .                                                               .
/// .                    Extension Field 2 (variable)               .
/// .                                                               .
/// |                                                               |
/// +--------------+---------------+---------------+----------------+
/// |                        Key Identifier                         |
/// +--------------+---------------+---------------+----------------+
/// |                                                               |
/// |                          dgst (128)                           |
/// |                                                               |
/// +--------------+---------------+---------------+----------------+
/// ```
#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
pub struct NtpPacket {
    pub header: Header,
    pub root_delay: ShortTimestamp,
    pub root_dispersion: ShortTimestamp,
    pub ref_id: RefIdentifier,
    pub ref_timestamp: Timestamp,
    pub origin_timestamp: Timestamp,
    pub receive_timestamp: Timestamp,
    pub transmit_timestamp: Timestamp,
}

impl BinLen for NtpPacket {
    const LEN: usize = Header::LEN
        + ShortTimestamp::LEN
        + ShortTimestamp::LEN
        + RefIdentifier::LEN
        + Timestamp::LEN
        + Timestamp::LEN
        + Timestamp::LEN
        + Timestamp::LEN;
}

impl ToNetPacket for NtpPacket {
    fn to_net(&self, v: &mut Vec<u8>) {
        self.header.to_net(v);
        self.root_delay.to_net(v);
        self.root_dispersion.to_net(v);
        self.ref_id.to_net(v);
        self.ref_timestamp.to_net(v);
        self.origin_timestamp.to_net(v);
        self.receive_timestamp.to_net(v);
        self.transmit_timestamp.to_net(v);
    }
}

impl NtpPacket {
    pub fn new() -> NtpPacket {
        NtpPacket {
            transmit_timestamp: Timestamp::now(),
            ..NtpPacket::default()
        }
    }
}