use crate::bigint::BigUint;
use crate::cert::build_min_cert;
use crate::keydata::{RSA_D, RSA_N, SPKI_DER};
use crate::sha256::sha256;

const APK_SIG_BLOCK_MAGIC: &[u8; 16] = b"APK Sig Block 42";
const V2_BLOCK_ID: u32 = 0x7109_871a;
const SIG_ALG_RSA_PKCS1_SHA256: u32 = 0x0103;

const CHUNK_SIZE: usize = 1024 * 1024;

fn v2_sig_len() -> usize {
    RSA_N.len()
}

pub fn add_v2_signature(zip: &[u8]) -> Vec<u8> {
    let eocd_pos = find_eocd(zip).expect("EOCD not found");
    let cd_offset = read_u32(zip, eocd_pos + 16) as usize;
    let cd_size = read_u32(zip, eocd_pos + 12) as usize;

    let section1 = &zip[0..cd_offset];
    let section2 = &zip[cd_offset..cd_offset + cd_size];
    let eocd = &zip[eocd_pos..];

    let sig_len = v2_sig_len();
    let dummy_digest = [0u8; 32];
    let dummy_block = build_signing_block(&dummy_digest, &vec![0u8; sig_len]);
    let block_len = dummy_block.len();

    let mut eocd_for_digest = eocd.to_vec();
    write_u32(&mut eocd_for_digest, 16, cd_offset as u32);

    let mut eocd_final = eocd.to_vec();
    write_u32(&mut eocd_final, 16, (cd_offset + block_len) as u32);

    let digest = chunked_sha256(&[section1, section2, &eocd_for_digest]);

    let (real_block, _) = build_signed_block(&digest);
    assert_eq!(
        real_block.len(),
        block_len,
        "signing block length changed between passes"
    );

    let mut out = Vec::with_capacity(zip.len() + block_len);
    out.extend_from_slice(section1);
    out.extend_from_slice(&real_block);
    out.extend_from_slice(section2);
    out.extend_from_slice(&eocd_final);
    out
}

fn build_signed_block(content_digest: &[u8; 32]) -> (Vec<u8>, Vec<u8>) {
    let signed_data = build_signed_data(content_digest);

    let signature = rsa_sign_sha256(&signed_data);

    let block = build_signing_block(content_digest, &signature);
    (block, signature)
}

fn build_signed_data(content_digest: &[u8]) -> Vec<u8> {
    let mut digest_entry = Vec::new();
    write_u32_push(&mut digest_entry, SIG_ALG_RSA_PKCS1_SHA256);
    push_len_prefixed(&mut digest_entry, content_digest);
    let mut digests = Vec::new();
    push_len_prefixed(&mut digests, &digest_entry);

    let cert = build_min_cert();
    let mut certs = Vec::new();
    push_len_prefixed(&mut certs, &cert);

    let mut sd = Vec::new();
    push_len_prefixed(&mut sd, &digests);
    push_len_prefixed(&mut sd, &certs);
    let additional: Vec<u8> = Vec::new();
    push_len_prefixed(&mut sd, &additional);
    sd
}

fn build_signing_block(content_digest: &[u8; 32], signature: &[u8]) -> Vec<u8> {
    let signed_data = build_signed_data(content_digest);

    let mut sig_entry = Vec::new();
    write_u32_push(&mut sig_entry, SIG_ALG_RSA_PKCS1_SHA256);
    push_len_prefixed(&mut sig_entry, signature);
    let mut signatures = Vec::new();
    push_len_prefixed(&mut signatures, &sig_entry);

    let mut signer = Vec::new();
    push_len_prefixed(&mut signer, &signed_data);
    push_len_prefixed(&mut signer, &signatures);
    push_len_prefixed(&mut signer, SPKI_DER);

    let mut signers = Vec::new();
    push_len_prefixed(&mut signers, &signer);

    let mut v2_value = Vec::new();
    push_len_prefixed(&mut v2_value, &signers);

    let pair_payload_len = 4 + v2_value.len();
    let mut pairs = Vec::new();
    write_u64_push(&mut pairs, pair_payload_len as u64);
    write_u32_push(&mut pairs, V2_BLOCK_ID);
    pairs.extend_from_slice(&v2_value);

    let size_of_block = (pairs.len() + 8 + 16) as u64;
    let mut block = Vec::new();
    write_u64_push(&mut block, size_of_block);
    block.extend_from_slice(&pairs);
    write_u64_push(&mut block, size_of_block);
    block.extend_from_slice(APK_SIG_BLOCK_MAGIC);
    block
}

fn chunked_sha256(sections: &[&[u8]]) -> [u8; 32] {
    let mut chunk_digests: Vec<[u8; 32]> = Vec::new();
    for sec in sections {
        let mut off = 0;
        while off < sec.len() {
            let end = (off + CHUNK_SIZE).min(sec.len());
            let chunk = &sec[off..end];
            let mut buf = Vec::with_capacity(1 + 4 + chunk.len());
            buf.push(0xa5);
            buf.extend_from_slice(&(chunk.len() as u32).to_le_bytes());
            buf.extend_from_slice(chunk);
            chunk_digests.push(sha256(&buf));
            off = end;
        }
    }
    let mut top = Vec::with_capacity(1 + 4 + chunk_digests.len() * 32);
    top.push(0x5a);
    top.extend_from_slice(&(chunk_digests.len() as u32).to_le_bytes());
    for d in &chunk_digests {
        top.extend_from_slice(d);
    }
    sha256(&top)
}

fn rsa_sign_sha256(message: &[u8]) -> Vec<u8> {
    let n = BigUint::from_be_bytes(RSA_N);
    let d = BigUint::from_be_bytes(RSA_D);
    let k = RSA_N.len();

    let digest = sha256(message);
    let sha256_prefix: &[u8] = &[
        0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01,
        0x05, 0x00, 0x04, 0x20,
    ];
    let mut t = Vec::new();
    t.extend_from_slice(sha256_prefix);
    t.extend_from_slice(&digest);

    let ps_len = k - t.len() - 3;
    let mut em = Vec::with_capacity(k);
    em.push(0x00);
    em.push(0x01);
    em.extend(std::iter::repeat(0xFF).take(ps_len));
    em.push(0x00);
    em.extend_from_slice(&t);

    let m = BigUint::from_be_bytes(&em);
    let s = m.modpow(&d, &n);
    s.to_be_bytes_padded(k)
}

fn push_len_prefixed(out: &mut Vec<u8>, data: &[u8]) {
    out.extend_from_slice(&(data.len() as u32).to_le_bytes());
    out.extend_from_slice(data);
}
fn write_u32_push(out: &mut Vec<u8>, x: u32) {
    out.extend_from_slice(&x.to_le_bytes());
}
fn write_u64_push(out: &mut Vec<u8>, x: u64) {
    out.extend_from_slice(&x.to_le_bytes());
}
fn read_u32(buf: &[u8], pos: usize) -> u32 {
    u32::from_le_bytes([buf[pos], buf[pos + 1], buf[pos + 2], buf[pos + 3]])
}
fn write_u32(buf: &mut [u8], pos: usize, x: u32) {
    buf[pos..pos + 4].copy_from_slice(&x.to_le_bytes());
}

fn find_eocd(buf: &[u8]) -> Option<usize> {
    if buf.len() < 22 {
        return None;
    }
    let mut pos = buf.len() - 22;
    loop {
        if read_u32(buf, pos) == 0x0605_4b50 {
            return Some(pos);
        }
        if pos == 0 {
            return None;
        }
        pos -= 1;
    }
}