// Copyright (C) 2024 Matthew Waters <matthew@centricular.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use stun_types::message::{
    IntegrityKey, Message, MessageHeader, MessageIntegrityCredentials, ShortTermCredentials,
};

static RFC5679_VECTOR1_DATA: [u8; 108] = [
    0x00, 0x01, 0x00, 0x58, // Request type message length
    0x21, 0x12, 0xa4, 0x42, // Magic cookie
    0xb7, 0xe7, 0xa7, 0x01, // }
    0xbc, 0x34, 0xd6, 0x86, // } Transaction ID
    0xfa, 0x87, 0xdf, 0xae, // }
    0x80, 0x22, 0x00, 0x10, // SOFTWARE header
    0x53, 0x54, 0x55, 0x4e, //   }
    0x20, 0x74, 0x65, 0x73, //   }  User-agent...
    0x74, 0x20, 0x63, 0x6c, //   }  ...name
    0x69, 0x65, 0x6e, 0x74, //   }
    0x00, 0x24, 0x00, 0x04, // PRIORITY header
    0x6e, 0x00, 0x01, 0xff, //   PRIORITY value
    0x80, 0x29, 0x00, 0x08, // ICE_CONTROLLED header
    0x93, 0x2f, 0xf9, 0xb1, //   Pseudo random number
    0x51, 0x26, 0x3b, 0x36, //   ... for tie breaker
    0x00, 0x06, 0x00, 0x09, // USERNAME header
    0x65, 0x76, 0x74, 0x6a, //   Username value
    0x3a, 0x68, 0x36, 0x76, //   (9 bytes)
    0x59, 0x20, 0x20, 0x20, //   (3 bytes padding)
    0x00, 0x08, 0x00, 0x14, // MESSAGE-INTEGRITY header
    0x9a, 0xea, 0xa7, 0x0c, //   }
    0xbf, 0xd8, 0xcb, 0x56, //   }
    0x78, 0x1e, 0xf2, 0xb5, //   } HMAC-SHA1 fingerprint
    0xb2, 0xd3, 0xf2, 0x49, //   }
    0xc1, 0xb5, 0x71, 0xa2, //   }
    0x80, 0x28, 0x00, 0x04, // FINGERPRINT header
    0xe5, 0x7a, 0x3b, 0xcf, //   CRC32 fingerprint
];
static RFC5679_VECTOR2_DATA: [u8; 80] = [
    0x01, 0x01, 0x00, 0x3c, // Response type message length
    0x21, 0x12, 0xa4, 0x42, // Magic cookie
    0xb7, 0xe7, 0xa7, 0x01, // }
    0xbc, 0x34, 0xd6, 0x86, // }  Transaction ID
    0xfa, 0x87, 0xdf, 0xae, // }
    0x80, 0x22, 0x00, 0x0b, // SOFTWARE attribute header
    0x74, 0x65, 0x73, 0x74, //   }
    0x20, 0x76, 0x65, 0x63, //   }  UTF-8 server name
    0x74, 0x6f, 0x72, 0x20, //   }
    0x00, 0x20, 0x00, 0x08, // XOR-MAPPED-ADDRESS attribute header
    0x00, 0x01, 0xa1, 0x47, //   Address family (IPv4) and xor'd mapped port number
    0xe1, 0x12, 0xa6, 0x43, //   Xor'd mapped IPv4 address
    0x00, 0x08, 0x00, 0x14, //   MESSAGE-INTEGRITY attribute header
    0x2b, 0x91, 0xf5, 0x99, // }
    0xfd, 0x9e, 0x90, 0xc3, // }
    0x8c, 0x74, 0x89, 0xf9, // }  HMAC-SHA1 fingerprint
    0x2a, 0xf9, 0xba, 0x53, // }
    0xf0, 0x6b, 0xe7, 0xd7, // }
    0x80, 0x28, 0x00, 0x04, //  FINGERPRINT attribute header
    0xc0, 0x7d, 0x4c, 0x96, //  CRC32 fingerprint
];

static RFC5679_VECTOR3_DATA: [u8; 92] = [
    0x01, 0x01, 0x00, 0x48, // Response type and message length
    0x21, 0x12, 0xa4, 0x42, // Magic cookie
    0xb7, 0xe7, 0xa7, 0x01, // }
    0xbc, 0x34, 0xd6, 0x86, // }  Transaction ID
    0xfa, 0x87, 0xdf, 0xae, // }
    0x80, 0x22, 0x00, 0x0b, //    SOFTWARE attribute header
    0x74, 0x65, 0x73, 0x74, // }
    0x20, 0x76, 0x65, 0x63, // }  UTF-8 server name
    0x74, 0x6f, 0x72, 0x20, // }
    0x00, 0x20, 0x00, 0x14, //    XOR-MAPPED-ADDRESS attribute header
    0x00, 0x02, 0xa1, 0x47, //    Address family (IPv6) and xor'd mapped port number
    0x01, 0x13, 0xa9, 0xfa, // }
    0xa5, 0xd3, 0xf1, 0x79, // }  Xor'd mapped IPv6 address
    0xbc, 0x25, 0xf4, 0xb5, // }
    0xbe, 0xd2, 0xb9, 0xd9, // }
    0x00, 0x08, 0x00, 0x14, //    MESSAGE-INTEGRITY attribute header
    0xa3, 0x82, 0x95, 0x4e, // }
    0x4b, 0xe6, 0x7b, 0xf1, // }
    0x17, 0x84, 0xc9, 0x7c, // }  HMAC-SHA1 fingerprint
    0x82, 0x92, 0xc2, 0x75, // }
    0xbf, 0xe3, 0xed, 0x41, // }
    0x80, 0x28, 0x00, 0x04, //    FINGERPRINT attribute header
    0xc8, 0xfb, 0x0b, 0x4c, //    CRC32 fingerprint
];

static RFC5679_VECTOR4_DATA: [u8; 116] = [
    0x00, 0x01, 0x00, 0x60, //    Request type and message length
    0x21, 0x12, 0xa4, 0x42, //    Magic cookie
    0x78, 0xad, 0x34, 0x33, // }
    0xc6, 0xad, 0x72, 0xc0, // }  Transaction ID
    0x29, 0xda, 0x41, 0x2e, // }
    0x00, 0x06, 0x00, 0x12, //    USERNAME attribute header
    0xe3, 0x83, 0x9e, 0xe3, // }
    0x83, 0x88, 0xe3, 0x83, // }
    0xaa, 0xe3, 0x83, 0x83, // }  Username value (18 bytes) and padding (2 bytes)
    0xe3, 0x82, 0xaf, 0xe3, // }
    0x82, 0xb9, 0x00, 0x00, // }
    0x00, 0x15, 0x00, 0x1c, //    NONCE attribute header
    0x66, 0x2f, 0x2f, 0x34, // }
    0x39, 0x39, 0x6b, 0x39, // }
    0x35, 0x34, 0x64, 0x36, // }
    0x4f, 0x4c, 0x33, 0x34, // }  Nonce value
    0x6f, 0x4c, 0x39, 0x46, // }
    0x53, 0x54, 0x76, 0x79, // }
    0x36, 0x34, 0x73, 0x41, // }
    0x00, 0x14, 0x00, 0x0b, //    REALM attribute header
    0x65, 0x78, 0x61, 0x6d, // }
    0x70, 0x6c, 0x65, 0x2e, // }  Realm value (11 bytes) and padding (1 byte)
    0x6f, 0x72, 0x67, 0x00, // }
    0x00, 0x08, 0x00, 0x14, //    MESSAGE-INTEGRITY attribute header
    0xf6, 0x70, 0x24, 0x65, // }
    0x6d, 0xd6, 0x4a, 0x3e, // }
    0x02, 0xb8, 0xe0, 0x71, // }  HMAC-SHA1 fingerprint
    0x2e, 0x85, 0xc9, 0xa2, // }
    0x8c, 0xa8, 0x96, 0x66, // }
];

static RFC8489_VECTOR1_DATA: [u8; 164] = [
    0x00, 0x01, 0x00, 0x90, //     Request type and message length
    0x21, 0x12, 0xa4, 0x42, //     Magic cookie
    0x78, 0xad, 0x34, 0x33, //  }
    0xc6, 0xad, 0x72, 0xc0, //  }  Transaction ID
    0x29, 0xda, 0x41, 0x2e, //  }
    0x00, 0x1e, 0x00, 0x20, //     USERHASH attribute header
    0x4a, 0x3c, 0xf3, 0x8f, //  }
    0xef, 0x69, 0x92, 0xbd, //  }
    0xa9, 0x52, 0xc6, 0x78, //  }
    0x04, 0x17, 0xda, 0x0f, //  }  Userhash value (32 bytes)
    0x24, 0x81, 0x94, 0x15, //  }
    0x56, 0x9e, 0x60, 0xb2, //  }
    0x05, 0xc4, 0x6e, 0x41, //  }
    0x40, 0x7f, 0x17, 0x04, //  }
    0x00, 0x15, 0x00, 0x29, //     NONCE attribute header
    0x6f, 0x62, 0x4d, 0x61, //  }
    0x74, 0x4a, 0x6f, 0x73, //  }
    0x32, 0x41, 0x41, 0x41, //  }
    0x43, 0x66, 0x2f, 0x2f, //  }
    0x34, 0x39, 0x39, 0x6b, //  }  Nonce value and padding (3 bytes)
    0x39, 0x35, 0x34, 0x64, //  }
    0x36, 0x4f, 0x4c, 0x33, //  }
    0x34, 0x6f, 0x4c, 0x39, //  }
    0x46, 0x53, 0x54, 0x76, //  }
    0x79, 0x36, 0x34, 0x73, //  }
    0x41, 0x00, 0x00, 0x00, //  }
    0x00, 0x14, 0x00, 0x0b, //     REALM attribute header
    0x65, 0x78, 0x61, 0x6d, //  }
    0x70, 0x6c, 0x65, 0x2e, //  }  Realm value (11 bytes) and padding (1 byte)
    0x6f, 0x72, 0x67, 0x00, //  }
    0x00, 0x1d, 0x00, 0x04, //    PASSWORD-ALGORITHM attribute header
    0x00, 0x02, 0x00, 0x00, //    PASSWORD-ALGORITHM value (4 bytes)
    0x00, 0x1c, 0x00, 0x20, //    MESSAGE-INTEGRITY-SHA256 attribute header
    0xb5, 0xc7, 0xbf, 0x00, // }
    0x5b, 0x6c, 0x52, 0xa2, // }
    0x1c, 0x51, 0xc5, 0xe8, // }
    0x92, 0xf8, 0x19, 0x24, // }  HMAC-SHA256 value
    0x13, 0x62, 0x96, 0xcb, // }
    0x92, 0x7c, 0x43, 0x14, // }
    0x93, 0x09, 0x27, 0x8c, // }
    0xc6, 0x51, 0x8e, 0x65, // }
];

fn parse(data: &[u8]) {
    let _msg = Message::from_bytes(data).unwrap();
}

fn parse_header(data: &[u8]) {
    let _msg = MessageHeader::from_bytes(data).unwrap();
}

fn validate_integrity(msg: &Message, credentials: &MessageIntegrityCredentials) {
    msg.validate_integrity(credentials).unwrap();
}

fn validate_integrity_with_key(msg: &Message, key: &IntegrityKey) {
    msg.validate_integrity_with_key(key).unwrap();
}

fn bench_message_parse(c: &mut Criterion) {
    let test_vectors = [
        (
            "RFC5679/Vector1",
            RFC5679_VECTOR1_DATA.as_slice(),
            Some(MessageIntegrityCredentials::ShortTerm(
                ShortTermCredentials::new("VOkJxbRl1RmTxUk/WvJxBt".to_owned()),
            )),
        ),
        (
            "RFC5679/Vector2",
            RFC5679_VECTOR2_DATA.as_slice(),
            Some(MessageIntegrityCredentials::ShortTerm(
                ShortTermCredentials::new("VOkJxbRl1RmTxUk/WvJxBt".to_owned()),
            )),
        ),
        (
            "RFC5679/Vector3",
            RFC5679_VECTOR3_DATA.as_slice(),
            Some(MessageIntegrityCredentials::ShortTerm(
                ShortTermCredentials::new("VOkJxbRl1RmTxUk/WvJxBt".to_owned()),
            )),
        ),
        (
            "RFC5679/Vector4",
            RFC5679_VECTOR4_DATA.as_slice(),
            // we cannot currently handle the message integrity for this
            None,
        ),
        (
            "RFC8489/Vector1",
            RFC8489_VECTOR1_DATA.as_slice(),
            // we cannot currently handle the message integrity for this
            None,
        ),
    ];

    let mut group = c.benchmark_group("Message/Parse");
    for test_vector in test_vectors {
        group.throughput(criterion::Throughput::Bytes(test_vector.1.len() as u64));
        group.bench_with_input(
            BenchmarkId::from_parameter(format!("{}/Header", test_vector.0)),
            test_vector.1,
            |b, data| b.iter(|| parse_header(data)),
        );
        group.bench_with_input(
            BenchmarkId::from_parameter(test_vector.0),
            test_vector.1,
            |b, data| b.iter(|| parse(data)),
        );
        if let Some(credentials) = test_vector.2 {
            let msg = Message::from_bytes(test_vector.1).unwrap();
            let key = credentials.make_key();
            group.bench_with_input(
                BenchmarkId::from_parameter(format!("{}/ValidateIntegrity", test_vector.0)),
                &(msg, credentials),
                |b, (msg, credentials)| b.iter(|| validate_integrity(msg, credentials)),
            );
            let msg = Message::from_bytes(test_vector.1).unwrap();
            group.bench_with_input(
                BenchmarkId::from_parameter(format!(
                    "{}/ValidateIntegrity/CachedKey",
                    test_vector.0
                )),
                &msg,
                |b, msg| b.iter(|| validate_integrity_with_key(msg, &key)),
            );
        }
    }
    group.finish();
}

criterion_group!(message_parse, bench_message_parse);
criterion_main!(message_parse);
