//! GIFT-256 fixsliced key schedule
//! Corresponds to gift256_key_schedule @ 0x80057559
//! Expands a 256-bit key into 120 u32 (480 bytes) of round keys

use super::sbox::sbox_bitsliced;
use super::interleave::{key_deinterleave_a, key_deinterleave_b, key_deinterleave_c, interleave_key_half};

/// NOT compensation table: offsets (u32 indices) that must be inverted at the end.
/// Corresponds to the mass of ks[offset] ^= 0xFFFFFFFF at the end of the decompilation.
const NOT_OFFSETS: [usize; 56] = [
    0x08, 0x09, 0x0D, 0x0E,
    0x10, 0x11, 0x15, 0x16,
    0x18, 0x19, 0x1D, 0x1E,
    0x20, 0x21, 0x25, 0x26,
    0x28, 0x29, 0x2D, 0x2E,
    0x30, 0x31, 0x35, 0x36,
    0x38, 0x39, 0x3D, 0x3E,
    0x40, 0x41, 0x45, 0x46,
    0x48, 0x49, 0x4D, 0x4E,
    0x50, 0x51, 0x55, 0x56,
    0x58, 0x59, 0x5D, 0x5E,
    0x60, 0x61, 0x65, 0x66,
    0x68, 0x69, 0x6D, 0x6E,
    0x70, 0x71, 0x75, 0x76,
];

/// GF extension function.
/// Decompiled:
///   uVar4 = (src >> param3 | src << (32-param3)) & 0x03030303 ^ dst;
///   result = uVar4 ^ (uVar4 & 0x03030303) << 6
///                   ^ (uVar4 & 0x0F0F0F0F) << 4
///                   ^ (uVar4 & 0x3F3F3F3F) << 2;
fn gf_expand(src: u32, dst: u32, rot: u32) -> u32 {
    let x = (src.rotate_right(rot) & 0x03030303) ^ dst;
    x ^ ((x & 0x03030303) << 6) ^ ((x & 0x0F0F0F0F) << 4) ^ ((x & 0x3F3F3F3F) << 2)
}

/// Key rotate: copy ks[base+0..7] forward to ks[base+8..15] with rotation.
/// Corresponds to gift256_key_rotate @ 0x80060a8f
///
/// After tracing the CFF state machine:
///   ks[base+15] = ks[base+7]
///   ks[base+14] = ks[base+6]
///   ...
///   ks[base+9]  = ks[base+1]
///   ks[base+8]  = ks[base+0]
fn key_rotate(ks: &mut [u32], base: usize) {
    for i in (0..8).rev() {
        ks[base + 8 + i] = ks[base + i];
    }
}

/// LFSR expansion: perform GF expansion on 8 key words.
/// Corresponds to gift256_key_lfsr_expand @ 0x8005a060
/// rot = 14 or 6 (rotation amount)
fn key_lfsr_expand(ks: &mut [u32], base: usize, rot: u32) {
    for i in 0..8 {
        let src_idx = base + i;
        let dst_idx = base + i - 16; // base is always >= 16
        ks[src_idx] = gf_expand(ks[src_idx], ks[dst_idx], rot);
    }
}

/// Full key schedule: expand 256-bit key into 120 u32 round keys.
pub fn key_schedule(key: &[u8; 32]) -> [u32; 120] {
    let mut ks = [0u32; 120];

    // === Step 1: Bit-interleave packing (State 3) ===
    // Read key as 8 u32 (little-endian)
    let mut kw = [0u32; 8];
    for i in 0..8 {
        kw[i] = u32::from_le_bytes([key[i * 4], key[i * 4 + 1], key[i * 4 + 2], key[i * 4 + 3]]);
    }

    // Interleave: two halves of 4 words each
    interleave_key_half(&kw[0..4], &mut ks[0..8]);
    interleave_key_half(&kw[4..8], &mut ks[8..16]);

    // === Step 2: 12-step round key expansion (State 4+5, 6 iterations) ===
    let mut round = 8usize;
    let mut rk_const_offset = 0x10usize; // u32 index

    for _step in 0..6 {
        // --- Sub-step A (State 4) ---
        key_rotate(&mut ks, round);
        round += 8;

        // S-Box on the newly written 8 words
        let mut sbox_buf: [u32; 8] = ks[round - 8..round].try_into().unwrap();
        sbox_bitsliced(&mut sbox_buf);
        ks[round - 8..round].copy_from_slice(&sbox_buf);

        // NOT compensation (4 specific offsets within this block)
        ks[round - 8] ^= 0xFFFFFFFF;
        ks[round - 7] ^= 0xFFFFFFFF;
        ks[round - 3] ^= 0xFFFFFFFF;
        ks[round - 2] ^= 0xFFFFFFFF;

        // Round constant injection
        ks[rk_const_offset] ^= 0xC000;

        // LFSR expand with rot=14
        key_lfsr_expand(&mut ks, round, 14);

        // --- Sub-step B (State 5) ---
        key_rotate(&mut ks, round);
        round += 8;

        let mut sbox_buf: [u32; 8] = ks[round - 8..round].try_into().unwrap();
        sbox_bitsliced(&mut sbox_buf);
        ks[round - 8..round].copy_from_slice(&sbox_buf);

        ks[round - 8] ^= 0xFFFFFFFF;
        ks[round - 7] ^= 0xFFFFFFFF;
        ks[round - 3] ^= 0xFFFFFFFF;
        ks[round - 2] ^= 0xFFFFFFFF;

        // LFSR expand with rot=6
        key_lfsr_expand(&mut ks, round, 6);

        rk_const_offset += 17; // iVar24 += 0x44 bytes = 17 u32 indices
    }

    // === Step 3: Post-processing (State 2 + State 1) ===
    // Deinterleave with 3 different mask patterns, each block covers 32 u32
    for block in 0..3 {
        let base = block * 32;

        // Type A: offsets +8..+15
        for j in 0..8 {
            ks[base + 8 + j] = key_deinterleave_a(ks[base + 8 + j]);
        }

        // Type B: offset +16
        ks[base + 16] = key_deinterleave_b(ks[base + 16]);

        // Type B: offsets +17..+23
        for j in 1..8 {
            ks[base + 16 + j] = key_deinterleave_b(ks[base + 16 + j]);
        }

        // Type C: offsets +24..+31
        for j in 0..8 {
            ks[base + 24 + j] = key_deinterleave_c(ks[base + 24 + j]);
        }
    }

    // Final round key deinterleave (ks[104..111] = offsets 0x1A0..0x1BC)
    for j in 0..8 {
        ks[104 + j] = key_deinterleave_a(ks[104 + j]);
    }

    // Initial NOT
    ks[8] ^= 0xFFFFFFFF;

    // Final NOT compensation
    for &off in &NOT_OFFSETS {
        ks[off] ^= 0xFFFFFFFF;
    }

    ks
}
