use anyhow::Result;
use blake2::{digest::Mac, Blake2s, Blake2sMac, Digest};
use byteorder::LittleEndian;
use rand::rngs::StdRng;
use rand::RngCore;
use rand::SeedableRng;
use snow::Builder;
use std::io;
use std::net::UdpSocket;
use x25519_dalek::PublicKey;
use x25519_dalek::StaticSecret;
use zerocopy::byteorder::U32;
use zerocopy::{AsBytes, FromBytes, FromZeroes};

const SAMPLE_SIZE: usize = 100_000;

const SIZE_TAG: usize = 16; // poly1305 tag
const SIZE_X25519_POINT: usize = 32; // x25519 public key
const SIZE_TIMESTAMP: usize = 12;

pub const TYPE_INITIATION: u32 = 1;
pub const TYPE_RESPONSE: u32 = 2;
pub const TYPE_COOKIE_REPLY: u32 = 3;

#[repr(packed)]
#[derive(Copy, Clone, FromZeroes, FromBytes, AsBytes)]
pub struct NoiseInitiation {
    pub f_type: U32<LittleEndian>,
    pub f_sender: U32<LittleEndian>,
    pub f_ephemeral: [u8; SIZE_X25519_POINT],
    pub f_static: [u8; SIZE_X25519_POINT + SIZE_TAG],
    pub f_timestamp: [u8; SIZE_TIMESTAMP + SIZE_TAG],
}

use blake2::digest::typenum::U16;
use serde::Serialize;
use std::time::{Instant, SystemTime, UNIX_EPOCH};

pub type TAI64N = [u8; 12];

const TAI64_EPOCH: u64 = 0x400000000000000a;

pub const ZERO: TAI64N = [0u8; 12];

pub fn now() -> TAI64N {
    // get system time as duration
    let sysnow = SystemTime::now();
    let delta = sysnow.duration_since(UNIX_EPOCH).unwrap();

    // convert to tai64n
    let tai64_secs = delta.as_secs() + TAI64_EPOCH;
    let tai64_nano = delta.subsec_nanos();

    // serialize
    let mut res = [0u8; 12];
    res[..8].copy_from_slice(&tai64_secs.to_be_bytes()[..]);
    res[8..].copy_from_slice(&tai64_nano.to_be_bytes()[..]);
    res
}

fn perform_handshake() -> Result<u128> {
    let private_key = StaticSecret::from([
        112, 3, 135, 100, 197, 196, 171, 109, 13, 243, 249, 30, 91, 137, 88, 174, 72, 94, 143, 122,
        2, 21, 227, 21, 222, 226, 127, 221, 253, 182, 120, 100,
    ]);
    let _ = PublicKey::from(&private_key);

    let remote_public_key = PublicKey::from([
        247, 61, 49, 57, 18, 173, 228, 54, 73, 122, 167, 57, 231, 254, 152, 22, 88, 209, 207, 114,
        175, 247, 76, 143, 22, 194, 15, 128, 120, 200, 145, 59,
    ]);

    let preshared_key: [u8; 32] = [0; 32];

    let mut initiator = Builder::new("Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s".parse()?)
        .local_private_key(&private_key.to_bytes())
        .remote_public_key(&remote_public_key.to_bytes())
        .psk(0, &preshared_key)
        .prologue("WireGuard v1 zx2c4 Jason@zx2c4.com".as_bytes())
        .build_initiator()?;

    let mut initiation: Vec<u8> = Vec::with_capacity(148);
    initiation.extend_from_slice(&[1, 0, 0, 0]);

    let mut id = [0u8; 4];
    let mut rng = StdRng::from_entropy();
    rng.fill_bytes(&mut id);
    initiation.extend_from_slice(&id);

    initiation.extend_from_slice(&[0u8; 140]);
    let timestamp = now();
    initiator.write_message(&timestamp, &mut initiation[8..116])?;

    let mut mac_key = Blake2s::<blake2::digest::typenum::U32>::new();
    mac_key.update("mac1----".as_bytes());
    mac_key.update(remote_public_key);
    let mac_key = mac_key.finalize();

    let mut hasher = <Blake2sMac<U16> as Mac>::new_from_slice(&mac_key)?;
    hasher.update(&initiation[..116]);
    initiation[116..132].copy_from_slice(&hasher.finalize().into_bytes());

    let socket = UdpSocket::bind("0.0.0.0:0")?;

    let start = Instant::now();

    let size = socket.send_to(&initiation, "127.0.0.1:31337")?;
    assert_eq!(size, initiation.len());

    let mut buf = [0u8; 92];
    let (number_of_bytes, _) = socket.recv_from(&mut buf)?;

    let end = Instant::now();

    assert_eq!(buf[0], 2);
    assert_eq!(number_of_bytes, 92);

    Ok(end.duration_since(start).as_nanos())
}

#[derive(Serialize)]
struct Record {
    time: u128,
}

fn main() -> Result<()> {
    let mut wtr = csv::Writer::from_writer(io::stdout());

    for _ in 0..SAMPLE_SIZE {
        wtr.serialize(Record {
            time: perform_handshake()?,
        })?;
    }

    wtr.flush()?;

    Ok(())
}