use rpki::data_model::roa::{IpPrefix, RoaAfi, RoaEContent};

fn len_bytes(len: usize) -> Vec<u8> {
    if len < 128 {
        vec![len as u8]
    } else {
        let mut tmp = Vec::new();
        let mut n = len;
        while n > 0 {
            tmp.push((n & 0xFF) as u8);
            n >>= 8;
        }
        tmp.reverse();
        let mut out = vec![0x80 | (tmp.len() as u8)];
        out.extend(tmp);
        out
    }
}

fn tlv(tag: u8, content: &[u8]) -> Vec<u8> {
    let mut out = vec![tag];
    out.extend(len_bytes(content.len()));
    out.extend_from_slice(content);
    out
}

fn der_integer_u64(v: u64) -> Vec<u8> {
    let mut bytes = Vec::new();
    let mut n = v;
    if n == 0 {
        bytes.push(0);
    } else {
        while n > 0 {
            bytes.push((n & 0xFF) as u8);
            n >>= 8;
        }
        bytes.reverse();
        if bytes[0] & 0x80 != 0 {
            bytes.insert(0, 0);
        }
    }
    tlv(0x02, &bytes)
}

fn der_octet_string(bytes: &[u8]) -> Vec<u8> {
    tlv(0x04, bytes)
}

fn der_bit_string_from_prefix(prefix_addr: &[u8], prefix_len: u16) -> Vec<u8> {
    let byte_len = ((prefix_len as usize) + 7) / 8;
    let unused = (byte_len * 8) as u16 - prefix_len;
    let mut bytes = prefix_addr[..byte_len.min(prefix_addr.len())].to_vec();
    while bytes.len() < byte_len {
        bytes.push(0);
    }
    if !bytes.is_empty() && unused != 0 {
        let mask = (1u8 << (unused as u8)) - 1;
        let last = bytes.len() - 1;
        bytes[last] &= !mask;
    }
    let mut content = vec![unused as u8];
    content.extend(bytes);
    tlv(0x03, &content)
}

fn der_sequence(children: Vec<Vec<u8>>) -> Vec<u8> {
    let mut content = Vec::new();
    for c in children {
        content.extend(c);
    }
    tlv(0x30, &content)
}

fn roa_ip_address(prefix_bs: Vec<u8>) -> Vec<u8> {
    der_sequence(vec![prefix_bs])
}

fn roa_ip_family(afi: [u8; 2], addresses: Vec<Vec<u8>>) -> Vec<u8> {
    der_sequence(vec![der_octet_string(&afi), der_sequence(addresses)])
}

fn roa_attestation(as_id: u64, families: Vec<Vec<u8>>) -> Vec<u8> {
    der_sequence(vec![der_integer_u64(as_id), der_sequence(families)])
}

#[test]
fn canonicalize_sorts_families_sorts_and_dedups_addresses() {
    // Build:
    // - families are in order: IPv6 then IPv4 (should become IPv4 then IPv6)
    // - IPv4 addresses contain a duplicate (should dedup)
    // - IPv4 /24 uses only 3 bytes, should be padded to 4 with host bits cleared
    let v6 = roa_ip_family(
        [0, 2],
        vec![roa_ip_address(der_bit_string_from_prefix(
            &[0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            32,
        ))],
    );

    let v4_prefix = der_bit_string_from_prefix(&[192, 0, 2, 255], 24);
    let v4 = roa_ip_family(
        [0, 1],
        vec![
            roa_ip_address(v4_prefix.clone()),
            roa_ip_address(v4_prefix), // duplicate
        ],
    );

    let der = roa_attestation(64496, vec![v6, v4]);
    let roa = RoaEContent::decode_der(&der).expect("decode roa econtent");

    assert_eq!(roa.ip_addr_blocks.len(), 2);
    assert_eq!(roa.ip_addr_blocks[0].afi, RoaAfi::Ipv4);
    assert_eq!(roa.ip_addr_blocks[1].afi, RoaAfi::Ipv6);

    let v4_fam = &roa.ip_addr_blocks[0];
    assert_eq!(v4_fam.addresses.len(), 1);
    assert_eq!(
        v4_fam.addresses[0].prefix,
        IpPrefix {
            afi: RoaAfi::Ipv4,
            prefix_len: 24,
            addr: vec![192, 0, 2, 0],
        }
    );
}