use super::*;
use hex;
use rand::rngs::OsRng;
use std::net::SocketAddr;
use std::thread;
use std::time::Duration;

use rand::prelude::*;

use x25519_dalek::PublicKey;
use x25519_dalek::StaticSecret;

use super::messages::{Initiation, Response};

fn setup_devices<R: RngCore + CryptoRng>(rng: &mut R) -> (PublicKey, Device, PublicKey, Device) {
    // generate new keypairs

    let sk1 = StaticSecret::new(rng);
    let pk1 = PublicKey::from(&sk1);

    let sk2 = StaticSecret::new(rng);
    let pk2 = PublicKey::from(&sk2);

    // pick random psk

    let mut psk = [0u8; 32];
    rng.fill_bytes(&mut psk[..]);

    // intialize devices on both ends

    let mut dev1 = Device::new();
    let mut dev2 = Device::new();

    dev1.set_sk(Some(sk1));
    dev2.set_sk(Some(sk2));

    dev1.add(pk2).unwrap();
    dev2.add(pk1).unwrap();

    dev1.set_psk(pk2, psk).unwrap();
    dev2.set_psk(pk1, psk).unwrap();

    (pk1, dev1, pk2, dev2)
}

fn wait() {
    thread::sleep(Duration::from_millis(20));
}

/* Test longest possible handshake interaction (7 messages):
 *
 * 1. I -> R (initation)
 * 2. I <- R (cookie reply)
 * 3. I -> R (initation)
 * 4. I <- R (response)
 * 5. I -> R (cookie reply)
 * 6. I -> R (initation)
 * 7. I <- R (response)
 */
#[test]
fn handshake_under_load() {
    let mut rng = OsRng::new().unwrap();
    let (_pk1, dev1, pk2, dev2) = setup_devices(&mut rng);

    let src1: SocketAddr = "172.16.0.1:8080".parse().unwrap();
    let src2: SocketAddr = "172.16.0.2:7070".parse().unwrap();

    // 1. device-1 : create first initation
    let msg_init = dev1.begin(&mut rng, &pk2).unwrap();

    // 2. device-2 : responds with CookieReply
    let msg_cookie = match dev2.process(&mut rng, &msg_init, Some(&src1)).unwrap() {
        (None, Some(msg), None) => msg,
        _ => panic!("unexpected response"),
    };

    // device-1 : processes CookieReply (no response)
    match dev1.process(&mut rng, &msg_cookie, Some(&src2)).unwrap() {
        (None, None, None) => (),
        _ => panic!("unexpected response"),
    }

    // avoid initation flood detection
    wait();

    // 3. device-1 : create second initation
    let msg_init = dev1.begin(&mut rng, &pk2).unwrap();

    // 4. device-2 : responds with noise response
    let msg_response = match dev2.process(&mut rng, &msg_init, Some(&src1)).unwrap() {
        (Some(_), Some(msg), Some(kp)) => {
            assert_eq!(kp.initiator, false);
            msg
        }
        _ => panic!("unexpected response"),
    };

    // 5. device-1 : responds with CookieReply
    let msg_cookie = match dev1.process(&mut rng, &msg_response, Some(&src2)).unwrap() {
        (None, Some(msg), None) => msg,
        _ => panic!("unexpected response"),
    };

    // device-2 : processes CookieReply (no response)
    match dev2.process(&mut rng, &msg_cookie, Some(&src1)).unwrap() {
        (None, None, None) => (),
        _ => panic!("unexpected response"),
    }

    // avoid initation flood detection
    wait();

    // 6. device-1 : create third initation
    let msg_init = dev1.begin(&mut rng, &pk2).unwrap();

    // 7. device-2 : responds with noise response
    let (msg_response, kp1) = match dev2.process(&mut rng, &msg_init, Some(&src1)).unwrap() {
        (Some(_), Some(msg), Some(kp)) => {
            assert_eq!(kp.initiator, false);
            (msg, kp)
        }
        _ => panic!("unexpected response"),
    };

    // device-1 : process noise response
    let kp2 = match dev1.process(&mut rng, &msg_response, Some(&src2)).unwrap() {
        (Some(_), None, Some(kp)) => {
            assert_eq!(kp.initiator, true);
            kp
        }
        _ => panic!("unexpected response"),
    };

    assert_eq!(kp1.send, kp2.recv);
    assert_eq!(kp1.recv, kp2.send);
}

#[test]
fn handshake_no_load() {
    let mut rng = OsRng::new().unwrap();
    let (pk1, mut dev1, pk2, mut dev2) = setup_devices(&mut rng);

    // do a few handshakes (every handshake should succeed)

    for i in 0..10 {
        println!("handshake : {}", i);

        // create initiation

        let msg1 = dev1.begin(&mut rng, &pk2).unwrap();

        println!("msg1 = {} : {} bytes", hex::encode(&msg1[..]), msg1.len());
        println!(
            "msg1 = {:?}",
            Initiation::parse(&msg1[..]).expect("failed to parse initiation")
        );

        // process initiation and create response

        let (_, msg2, ks_r) = dev2
            .process(&mut rng, &msg1, None)
            .expect("failed to process initiation");

        let ks_r = ks_r.unwrap();
        let msg2 = msg2.unwrap();

        println!("msg2 = {} : {} bytes", hex::encode(&msg2[..]), msg2.len());
        println!(
            "msg2 = {:?}",
            Response::parse(&msg2[..]).expect("failed to parse response")
        );

        assert!(!ks_r.initiator, "Responders key-pair is confirmed");

        // process response and obtain confirmed key-pair

        let (_, msg3, ks_i) = dev1
            .process(&mut rng, &msg2, None)
            .expect("failed to process response");
        let ks_i = ks_i.unwrap();

        assert!(msg3.is_none(), "Returned message after response");
        assert!(ks_i.initiator, "Initiators key-pair is not confirmed");

        assert_eq!(ks_i.send, ks_r.recv, "KeyI.send != KeyR.recv");
        assert_eq!(ks_i.recv, ks_r.send, "KeyI.recv != KeyR.send");

        dev1.release(ks_i.local_id());
        dev2.release(ks_r.local_id());

        // avoid initation flood detection
        wait();
    }

    dev1.remove(pk2).unwrap();
    dev2.remove(pk1).unwrap();
}