use std::fs;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::path::{Path, PathBuf};

use anyhow::{Context, Result};
use sha2::{Digest, Sha256};

const CONTENT_TYPE_OID: &str = "1.2.840.113549.1.9.16.1.54";
const SHA256_OID: &str = "2.16.840.1.101.3.4.2.1";

#[derive(Clone)]
struct Vrp {
    addr: IpAddr,
    prefix_len: u8,
    max_len: u8,
    asn: u32,
}

#[derive(Clone)]
struct Vap {
    customer_asn: u32,
    providers: Vec<u32>,
}

fn main() -> Result<()> {
    let out_dir = parse_out_dir_arg();
    fs::create_dir_all(&out_dir)
        .with_context(|| format!("failed to create output dir {}", out_dir.display()))?;

    write_snapshot(
        &out_dir.join("20260403T000001Z-mini-a.ccr"),
        "20260403000001Z",
        vec![
            Vrp {
                addr: IpAddr::V4(Ipv4Addr::new(10, 0, 0, 0)),
                prefix_len: 24,
                max_len: 24,
                asn: 65001,
            },
            Vrp {
                addr: IpAddr::V6(Ipv6Addr::new(0x2001, 0xdb8, 1, 0, 0, 0, 0, 0)),
                prefix_len: 48,
                max_len: 48,
                asn: 65002,
            },
        ],
        vec![Vap {
            customer_asn: 65010,
            providers: vec![65011, 65012],
        }],
    )?;

    write_snapshot(
        &out_dir.join("20260403T000101Z-mini-b.ccr"),
        "20260403000101Z",
        vec![
            Vrp {
                addr: IpAddr::V4(Ipv4Addr::new(10, 0, 0, 0)),
                prefix_len: 24,
                max_len: 24,
                asn: 65001,
            },
            Vrp {
                addr: IpAddr::V4(Ipv4Addr::new(10, 0, 1, 0)),
                prefix_len: 24,
                max_len: 24,
                asn: 65003,
            },
            Vrp {
                addr: IpAddr::V6(Ipv6Addr::new(0x2001, 0xdb8, 1, 0, 0, 0, 0, 0)),
                prefix_len: 48,
                max_len: 48,
                asn: 65002,
            },
        ],
        vec![Vap {
            customer_asn: 65010,
            providers: vec![65011, 65012, 65013],
        }],
    )?;

    write_snapshot(
        &out_dir.join("20260403T000201Z-mini-c.ccr"),
        "20260403000201Z",
        vec![
            Vrp {
                addr: IpAddr::V4(Ipv4Addr::new(10, 0, 1, 0)),
                prefix_len: 24,
                max_len: 24,
                asn: 65003,
            },
            Vrp {
                addr: IpAddr::V6(Ipv6Addr::new(0x2001, 0xdb8, 2, 0, 0, 0, 0, 0)),
                prefix_len: 48,
                max_len: 48,
                asn: 65004,
            },
        ],
        vec![
            Vap {
                customer_asn: 65010,
                providers: vec![65012, 65013],
            },
            Vap {
                customer_asn: 65020,
                providers: vec![65021],
            },
        ],
    )?;

    println!("generated CCR fixtures under {}", out_dir.display());
    Ok(())
}

fn parse_out_dir_arg() -> PathBuf {
    let mut args = std::env::args().skip(1);
    let mut out_dir = PathBuf::from("data");
    while let Some(arg) = args.next() {
        if arg == "--out-dir" {
            if let Some(v) = args.next() {
                out_dir = PathBuf::from(v);
            }
        }
    }
    out_dir
}

fn write_snapshot(path: &Path, produced_at: &str, vrps: Vec<Vrp>, vaps: Vec<Vap>) -> Result<()> {
    let bytes = encode_ccr_snapshot(produced_at, vrps, vaps);
    fs::write(path, bytes).with_context(|| format!("failed to write {}", path.display()))?;
    println!("wrote {}", path.display());
    Ok(())
}

fn encode_ccr_snapshot(produced_at: &str, vrps: Vec<Vrp>, vaps: Vec<Vap>) -> Vec<u8> {
    let vrp_sets = vrps
        .into_iter()
        .map(encode_roa_payload_set)
        .collect::<Vec<_>>();
    let vap_sets = vaps
        .into_iter()
        .map(encode_aspa_payload_set)
        .collect::<Vec<_>>();

    let vrp_set_seq = der_sequence(vrp_sets);
    let vap_set_seq = der_sequence(vap_sets);

    let vrp_hash = Sha256::digest(&vrp_set_seq).to_vec();
    let vap_hash = Sha256::digest(&vap_set_seq).to_vec();

    // draft-ietf-sidrops-rpki-ccr-02:
    // ROAPayloadState/ASPAPayloadState include payload-set sequence + hash.
    let vrp_state = der_sequence(vec![vrp_set_seq, der_octet_string(vrp_hash)]);
    let vap_state = der_sequence(vec![vap_set_seq, der_octet_string(vap_hash)]);

    // AlgorithmIdentifier for SHA-256.
    let hash_alg = der_sequence(vec![der_oid(SHA256_OID), der_null()]);

    let payload = der_sequence(vec![
        der_integer(0),
        hash_alg,
        der_generalized_time(produced_at),
        der_ctx(2, vrp_state),
        der_ctx(3, vap_state),
    ]);

    der_sequence(vec![der_oid(CONTENT_TYPE_OID), der_ctx(0, payload)])
}

fn encode_roa_payload_set(v: Vrp) -> Vec<u8> {
    let (afi, addr_bytes) = match v.addr {
        IpAddr::V4(ip) => ([0u8, 1u8].to_vec(), ip.octets().to_vec()),
        IpAddr::V6(ip) => ([0u8, 2u8].to_vec(), ip.octets().to_vec()),
    };
    let bit_string = prefix_to_bit_string(&addr_bytes, v.prefix_len);
    let roa_ip = der_sequence(vec![
        der_bit_string(0, bit_string),
        der_integer(u32::from(v.max_len)),
    ]);
    let family = der_sequence(vec![der_octet_string(afi), der_sequence(vec![roa_ip])]);
    der_sequence(vec![der_integer(v.asn), der_sequence(vec![family])])
}

fn encode_aspa_payload_set(v: Vap) -> Vec<u8> {
    let providers = v.providers.into_iter().map(der_integer).collect::<Vec<_>>();
    der_sequence(vec![der_integer(v.customer_asn), der_sequence(providers)])
}

fn prefix_to_bit_string(addr: &[u8], prefix_len: u8) -> Vec<u8> {
    let byte_len = usize::from(prefix_len).div_ceil(8);
    let mut out = addr[..byte_len].to_vec();
    let rem = prefix_len % 8;
    if rem != 0 {
        let mask = 0xFFu8 << (8 - rem);
        let last = out.len() - 1;
        out[last] &= mask;
    }
    out
}

fn der_sequence(items: Vec<Vec<u8>>) -> Vec<u8> {
    let content = items.concat();
    der_tlv(0x30, content)
}

fn der_integer(v: u32) -> Vec<u8> {
    if v == 0 {
        return der_tlv(0x02, vec![0]);
    }
    let mut bytes = v.to_be_bytes().to_vec();
    while bytes.len() > 1 && bytes[0] == 0 {
        bytes.remove(0);
    }
    if bytes[0] & 0x80 != 0 {
        bytes.insert(0, 0);
    }
    der_tlv(0x02, bytes)
}

fn der_oid(oid: &str) -> Vec<u8> {
    let parts = oid
        .split('.')
        .map(|s| s.parse::<u32>().unwrap())
        .collect::<Vec<_>>();
    assert!(parts.len() >= 2);
    let mut out = Vec::new();
    out.push((parts[0] * 40 + parts[1]) as u8);
    for &part in &parts[2..] {
        out.extend(base128(part));
    }
    der_tlv(0x06, out)
}

fn base128(mut n: u32) -> Vec<u8> {
    let mut buf = vec![(n & 0x7F) as u8];
    n >>= 7;
    while n > 0 {
        buf.push(((n & 0x7F) as u8) | 0x80);
        n >>= 7;
    }
    buf.reverse();
    buf
}

fn der_octet_string(bytes: Vec<u8>) -> Vec<u8> {
    der_tlv(0x04, bytes)
}

fn der_null() -> Vec<u8> {
    der_tlv(0x05, Vec::new())
}

fn der_bit_string(unused_bits: u8, bytes: Vec<u8>) -> Vec<u8> {
    let mut content = Vec::with_capacity(1 + bytes.len());
    content.push(unused_bits);
    content.extend(bytes);
    der_tlv(0x03, content)
}

fn der_generalized_time(v: &str) -> Vec<u8> {
    der_tlv(0x18, v.as_bytes().to_vec())
}

fn der_ctx(tag_no: u8, encoded_inner_der: Vec<u8>) -> Vec<u8> {
    der_tlv(0xA0 + tag_no, encoded_inner_der)
}

fn der_tlv(tag: u8, content: Vec<u8>) -> Vec<u8> {
    let mut out = Vec::with_capacity(2 + content.len());
    out.push(tag);
    out.extend(der_len(content.len()));
    out.extend(content);
    out
}

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