use crate::keys::DnssecKeyPair;
use crate::keytag::compute_key_tag;
use rustdns_protocol::name::encode_name;
use rustdns_protocol::packet::{encode_rrsig, DnsRecord};
use rustdns_protocol::types::QType;
use sha2::{Sha256, Digest};

/// Canonical RRset serialization for DNSSEC signing (RFC 4034 Section 6).
/// Each record: name(wire) + type(2) + class(2) + ttl(4) + rdlength(2) + rdata
pub fn serialize_rrset_canonical(records: &[DnsRecord]) -> Vec<u8> {
    let mut buf = Vec::new();
    for rr in records {
        if rr.rtype == QType::OPT {
            continue;
        }
        let name = if rr.name.ends_with('.') {
            rr.name.to_lowercase()
        } else {
            format!("{}.", rr.name).to_lowercase()
        };
        buf.extend_from_slice(&encode_name(&name));
        buf.extend_from_slice(&rr.rtype.to_u16().to_be_bytes());
        buf.extend_from_slice(&rr.rclass.to_u16().to_be_bytes());
        buf.extend_from_slice(&rr.ttl.to_be_bytes());
        buf.extend_from_slice(&(rr.rdata.len() as u16).to_be_bytes());
        buf.extend_from_slice(&rr.rdata);
    }
    buf
}

/// Generate an RRSIG record for a given RRset.
pub fn generate_rrsig(
    key_pair: &DnssecKeyPair,
    zone: &str,
    rrset: &[DnsRecord],
    name: &str,
    rtype: QType,
) -> DnsRecord {
    let algorithm = key_pair.algorithm().number();
    let dnskey_rdata = key_pair.dnskey_rdata();
    let key_tag = compute_key_tag(&dnskey_rdata);

    let signers_name = if zone.ends_with('.') {
        zone.to_string()
    } else {
        format!("{}.", zone)
    };

    let ttl = if rrset.is_empty() { 300 } else { rrset[0].ttl };
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_secs() as u32;
    let inception = now.wrapping_sub(3600); // 1 hour ago
    let expiration = inception.wrapping_add(86400); // +1 day

    let labels = name
        .strip_suffix('.')
        .unwrap_or(name)
        .split('.')
        .filter(|l| !l.is_empty())
        .count() as u8;

    // Build the RRSIG RDATA preamble (everything before the signature)
    let type_covered = rtype.to_u16();
    let mut sig_data = Vec::new();
    sig_data.extend_from_slice(&type_covered.to_be_bytes());
    sig_data.push(algorithm);
    sig_data.push(labels);
    sig_data.extend_from_slice(&ttl.to_be_bytes());
    sig_data.extend_from_slice(&expiration.to_be_bytes());
    sig_data.extend_from_slice(&inception.to_be_bytes());
    sig_data.extend_from_slice(&key_tag.to_be_bytes());
    sig_data.extend_from_slice(&encode_name(&signers_name));

    // Append the canonical RRset
    sig_data.extend_from_slice(&serialize_rrset_canonical(rrset));

    // Sign: ECDSA uses SHA-256 internally via the p256 crate, ED25519 does its own hashing
    let signature = match key_pair {
        DnssecKeyPair::EcdsaP256 { .. } => {
            // For ECDSA, we hash first then sign
            let hash = Sha256::digest(&sig_data);
            key_pair.sign(&hash)
        }
        DnssecKeyPair::Ed25519 { .. } => {
            // ED25519 includes hashing internally
            key_pair.sign(&sig_data)
        }
    };

    let rrsig_rdata = encode_rrsig(
        type_covered,
        algorithm,
        labels,
        ttl,
        expiration,
        inception,
        key_tag,
        &signers_name,
        &signature,
    );

    DnsRecord {
        name: name.to_string(),
        rtype: QType::RRSIG,
        rclass: rustdns_protocol::types::QClass::IN,
        ttl,
        rdata: rrsig_rdata,
        opt_flags: None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::keys::{DnssecAlgorithm, DnssecKeyPair};
    use rustdns_protocol::packet::{build_record, encode_a};

    #[test]
    fn test_generate_rrsig_ecdsa() {
        let kp = DnssecKeyPair::generate(DnssecAlgorithm::EcdsaP256Sha256);
        let record = build_record("test.example.com", QType::A, 300, encode_a("127.0.0.1"));
        let rrsig = generate_rrsig(&kp, "example.com", &[record], "test.example.com", QType::A);

        assert_eq!(rrsig.rtype, QType::RRSIG);
        assert!(!rrsig.rdata.is_empty());
    }

    #[test]
    fn test_generate_rrsig_ed25519() {
        let kp = DnssecKeyPair::generate(DnssecAlgorithm::Ed25519);
        let record = build_record("test.example.com", QType::A, 300, encode_a("10.0.0.1"));
        let rrsig = generate_rrsig(&kp, "example.com", &[record], "test.example.com", QType::A);

        assert_eq!(rrsig.rtype, QType::RRSIG);
        assert!(!rrsig.rdata.is_empty());
    }

    #[test]
    fn test_serialize_rrset_canonical() {
        let r1 = build_record("example.com", QType::A, 300, encode_a("1.2.3.4"));
        let r2 = build_record("example.com", QType::A, 300, encode_a("5.6.7.8"));
        let serialized = serialize_rrset_canonical(&[r1, r2]);
        assert!(!serialized.is_empty());
    }
}