package bangcle

// White-box AES block functions.
//
// The index arithmetic below is a one-for-one transcription of decompiled code
// (via pybyd). It is deliberately not tidied or renamed: keeping it diffable
// against the reference is worth more than readable variable names, because the
// only way to debug a mistake here is to compare intermediate state with pybyd.

import "encoding/binary"

// prepareMatrix transposes a 4x4 block into the working state layout (col*8+row).
func prepareMatrix(in, state []byte) {
	for col := 0; col < 4; col++ {
		for row := 0; row < 4; row++ {
			state[col*8+row] = in[col+row*4]
		}
	}
}

func decryptBlock(t *lookupTables, block []byte) []byte {
	state := make([]byte, 32)
	temp64 := make([]byte, 64)
	tmp32 := make([]byte, 32)
	out := make([]byte, 16)

	prepareMatrix(block, state)

	for rnd := 9; rnd >= 1; rnd-- {
		l21 := rnd * 4
		permPtr := 0
		for i := 0; i < 4; i++ {
			b := int(t.permDecrypt[permPtr])
			l16 := i * 8
			base := i * 16
			for j := 0; j < 4; j++ {
				u7 := (b + j) & 3
				idx := int(state[l16+u7]) + (i+(l21+u7)*4)*256
				binary.LittleEndian.PutUint32(temp64[base+j*4:], binary.LittleEndian.Uint32(t.invRound[idx*4:]))
			}
			permPtr += 2
		}

		i15 := 1
		for l21x := 0; l21x < 4; l21x++ {
			pb := l21x
			for l9 := 0; l9 < 4; l9++ {
				u6 := int(temp64[pb]) & 0xF
				u26 := int(temp64[pb]) & 0xF0
				f := [3]int{int(temp64[pb+0x10]), int(temp64[pb+0x20]), int(temp64[pb+0x30])}

				l2 := l9*0x18 + rnd*0x60
				i25 := i15
				for k := 0; k < 3; k++ {
					b := f[k]
					u27 := u6 | ((b << 4) & 0xFF)
					u26 = ((u26 >> 4) | ((b >> 4) << 4)) & 0xFF
					u6 = int(t.invXor[(l2+i25-1)*0x100+u27]) & 0xF
					u26 = (int(t.invXor[(l2+i25)*0x100+u26]) & 0xF) << 4
					i25 += 2
				}
				state[l9+l21x*8] = byte((u26 | u6) & 0xFF)
				pb += 4
			}
			i15 += 6
		}
	}

	copy(tmp32, state)
	u8, u10, u12 := 1, 3, 2
	for row := 0; row < 4; row++ {
		state[row] = t.invFirst[int(tmp32[row])+row*0x400]

		r1 := u10 & 3
		state[8+row] = t.invFirst[int(tmp32[8+r1])+r1*0x400+0x100]

		r2 := u12 & 3
		state[0x10+row] = t.invFirst[int(tmp32[0x10+r2])+r2*0x400+0x200]

		r3 := u8 & 3
		state[0x18+row] = t.invFirst[int(tmp32[0x18+r3])+r3*0x400+0x300]

		u8++
		u10++
		u12++
	}

	for col := 0; col < 4; col++ {
		for row := 0; row < 4; row++ {
			out[col+row*4] = state[col*8+row]
		}
	}
	return out
}

func encryptBlock(t *lookupTables, block []byte) []byte {
	state := make([]byte, 32)
	temp64 := make([]byte, 64)
	tmp32 := make([]byte, 32)
	out := make([]byte, 16)

	prepareMatrix(block, state)

	for rnd := 0; rnd < 9; rnd++ {
		l21 := rnd * 4
		permPtr := 0
		for i := 0; i < 4; i++ {
			b := int(t.permEncrypt[permPtr])
			l16 := i * 8
			base := i * 16
			for j := 0; j < 4; j++ {
				u8 := (b + j) & 3
				idx := int(state[l16+u8]) + (i+(l21+u8)*4)*256
				binary.LittleEndian.PutUint32(temp64[base+j*4:], binary.LittleEndian.Uint32(t.round[idx*4:]))
			}
			permPtr += 2
		}

		i16 := 1
		for l22 := 0; l22 < 4; l22++ {
			pb := l22
			for l10 := 0; l10 < 4; l10++ {
				u7 := int(temp64[pb]) & 0xF
				u26 := int(temp64[pb]) & 0xF0
				f := [3]int{int(temp64[pb+0x10]), int(temp64[pb+0x20]), int(temp64[pb+0x30])}

				l2 := l10*0x18 + rnd*0x60
				i25 := i16
				for k := 0; k < 3; k++ {
					b := f[k]
					u27 := u7 | ((b << 4) & 0xFF)
					u26 = ((u26 >> 4) | ((b >> 4) << 4)) & 0xFF
					u7 = int(t.xor[(l2+i25-1)*0x100+u27]) & 0xF
					u26 = (int(t.xor[(l2+i25)*0x100+u26]) & 0xF) << 4
					i25 += 2
				}
				state[l10+l22*8] = byte((u26 | u7) & 0xFF)
				pb += 4
			}
			i16 += 6
		}
	}

	copy(tmp32, state)
	for row := 0; row < 4; row++ {
		r0 := row & 3
		state[row] = t.final[int(tmp32[r0])+r0*0x400]

		r1 := (1 + row) & 3
		state[8+row] = t.final[int(tmp32[8+r1])+r1*0x400+0x100]

		r2 := (2 + row) & 3
		state[0x10+row] = t.final[int(tmp32[0x10+r2])+r2*0x400+0x200]

		r3 := (3 + row) & 3
		state[0x18+row] = t.final[int(tmp32[0x18+r3])+r3*0x400+0x300]
	}

	for col := 0; col < 4; col++ {
		for row := 0; row < 4; row++ {
			out[col+row*4] = state[col*8+row]
		}
	}
	return out
}