// Copyright 2025 International Digital Economy Academy
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

///|
fn flipWord(word : UInt) -> UInt {
  let b0 = (word & 0xff) << 24
  let b1 = (word & 0xff00) << 8
  let b2 = (word & 0xff0000) >> 8
  let b3 = (word & 0xff000000) >> 24
  b0 | b1 | b2 | b3
}

///|
test "flipWord" {
  let word = 0x12345678U
  let flipped = flipWord(word)
  inspect(flipped, content="2018915346")
}

///|
fn FixedArray::quarterRound(
  state : FixedArray[UInt],
  w : Int,
  x : Int,
  y : Int,
  z : Int,
) -> Unit {
  let mut a = state[w]
  let mut b = state[x]
  let mut c = state[y]
  let mut d = state[z]
  a += b
  d = d ^ a
  d = rotate_left_u(d, 16)
  c += d
  b = b ^ c
  b = rotate_left_u(b, 12)
  a += b
  d = d ^ a
  d = rotate_left_u(d, 8)
  c += d
  b = b ^ c
  b = rotate_left_u(b, 7)
  state[w] = a
  state[x] = b
  state[y] = c
  state[z] = d
}

///|
test "quarterRound" {
  let state = FixedArray::make(16, 0U)
  state[0] = 0x879531e0U
  state[1] = 0xc5ecf37dU
  state[2] = 0x516461b1U
  state[3] = 0xc9a62f8aU
  state[4] = 0x44c20ef3U
  state[5] = 0x3390af7fU
  state[6] = 0xd9fc690bU
  state[7] = 0x2a5f714cU
  state[8] = 0x53372767U
  state[9] = 0xb00a5631U
  state[10] = 0x974c541aU
  state[11] = 0x359e9963U
  state[12] = 0x5c971061U
  state[13] = 0x3d631689U
  state[14] = 0x2098d9d6U
  state[15] = 0x91dbd320U
  state.quarterRound(2, 7, 8, 13)
  debug_inspect(
    Array::from_iter(state.iter()),
    content=(
      #|[
      #|  2274701792,
      #|  3320640381,
      #|  3182986972,
      #|  3383111562,
      #|  1153568499,
      #|  865120127,
      #|  3657197835,
      #|  3484200914,
      #|  3832277632,
      #|  2953467441,
      #|  2538361882,
      #|  899586403,
      #|  1553404001,
      #|  3435166841,
      #|  546888150,
      #|  2447102752,
      #|]
    ),
  )
}

///|
fn FixedArray::chachaBlockRound(state : FixedArray[UInt]) -> Unit {
  state
  ..quarterRound(0, 4, 8, 12)
  ..quarterRound(1, 5, 9, 13)
  ..quarterRound(2, 6, 10, 14)
  ..quarterRound(3, 7, 11, 15)
  ..quarterRound(0, 5, 10, 15)
  ..quarterRound(1, 6, 11, 12)
  ..quarterRound(2, 7, 8, 13)
  .quarterRound(3, 4, 9, 14)
}

///|
test "chachaBlockRound" {
  let state = FixedArray::make(16, 0U)
  state[0] = 0x61707865U
  state[1] = 0x3320646eU
  state[2] = 0x79622d32U
  state[3] = 0x6b206574U
  state[4] = 0x03020100U
  state[5] = 0x07060504U
  state[6] = 0x0b0a0908U
  state[7] = 0x0f0e0d0cU
  state[8] = 0x13121110U
  state[9] = 0x17161514U
  state[10] = 0x1b1a1918U
  state[11] = 0x1f1e1d1cU
  state[12] = 0x00000001U
  state[13] = 0x00000000U
  state[14] = 0x00000000U
  state[15] = 0x00000000U
  state.chachaBlockRound()
  debug_inspect(
    Array::from_iter(state.iter()),
    content=(
      #|[
      #|  986087425,
      #|  3489031050,
      #|  2890662805,
      #|  2683391196,
      #|  1720476390,
      #|  1116253759,
      #|  2262580386,
      #|  3212003942,
      #|  2202368212,
      #|  756352536,
      #|  496298475,
      #|  669838588,
      #|  567302638,
      #|  1860562437,
      #|  1434237441,
      #|  2097484794,
      #|]
    ),
  )
}

///|
fn FixedArray::chachaBlockLoop(state : FixedArray[UInt], n : UInt) -> Unit {
  for _ in 0U.. FixedArray[Byte] raise Error {
  guard key.length() == 8 else {
    fail("Invalid key length -- key must be 8 32-bit unsigned integers")
  }
  chacha(key.map(flipWord), counter, block, 4, [nonce, nonce, nonce])
}

///|
/// Encrypts a block of data using the ChaCha12 algorithm.
/// - [key] must be 8 32-bit unsigned integers.
/// - [counter] is the counter value.
/// - [block] is the block of data to be encrypted.
/// - [nonce] is default to 0
/// - Returns the encrypted block of data.
#deprecated("Use ChaCha::chacha12 and ChaCha::transform instead")
pub fn[Data : ByteSource] chacha12(
  key : FixedArray[UInt],
  counter : UInt,
  block : Data,
  nonce? : UInt = 0,
) -> FixedArray[Byte] raise Error {
  guard key.length() == 8 else {
    fail("Invalid key length -- key must be 8 32-bit unsigned integers")
  }
  chacha(key.map(flipWord), counter, block, 6, [nonce, nonce, nonce])
}

///|
/// Encrypts a block of data using the ChaCha20 algorithm.
/// - [key] must be 8 32-bit unsigned integers.
/// - [counter] is the counter value.
/// - [block] is the block of data to be encrypted.
/// - [nonce] is default to 0
/// - Returns the encrypted block of data.
#deprecated("Use ChaCha::chacha20 and ChaCha::transform instead")
pub fn[Data : ByteSource] chacha20(
  key : FixedArray[UInt],
  counter : UInt,
  block : Data,
  nonce? : UInt = 0,
) -> FixedArray[Byte] raise Error {
  guard key.length() == 8 else {
    fail("Invalid key length -- key must be 8 32-bit unsigned integers")
  }
  chacha(key.map(flipWord), counter, block, 10, [nonce, nonce, nonce])
}

///|
#coverage.skip
fn[Data : ByteSource] chacha(
  key : FixedArray[UInt],
  counter : UInt,
  block : Data,
  round : UInt,
  nonce : FixedArray[UInt],
) -> FixedArray[Byte] {
  let block_length = block.length()
  guard block_length > 0 else { FixedArray::make(0, Byte::default()) }
  let block_count = (block_length - 1) / 64 + 1
  let available_blocks = 0x100000000UL - counter.to_uint64()
  guard block_count.to_uint64() <= available_blocks else {
    abort("ChaCha counter exhausted")
  }
  let buffer = FixedArray::make(block_length, Byte::default())
  let key_stream = FixedArray::make(64, b'\x00')
  for i = 0; i < block_length; i = i + 64 {
    chachaBlock(
      key,
      counter + i.reinterpret_as_uint() / 64,
      nonce,
      round,
      key_stream,
    )
    let len = @cmp.minimum(block_length - i, 64)
    chacha_xor_into(block, i, buffer, i, key_stream, 0, len)
  }
  buffer
}

///|
/// Errors raised by ChaCha operations.
pub suberror ChaChaError {
  CounterExhausted
  InvalidOutputRange
}

///|
struct ChaCha {
  key : FixedArray[UInt]
  nonce : FixedArray[UInt]
  key_stream : FixedArray[Byte]
  mut counter : UInt
  mut exhausted : Bool
  mut offset : Int
  round : UInt
}

///|
fn[K : ByteSource, N : ByteSource] ChaCha::new(
  key : K,
  nonce : N,
  round : UInt,
  counter : UInt,
) -> ChaCha raise Error {
  guard key.length() == 32 else {
    fail("Invalid key length -- key must be 256 bits")
  }
  guard nonce.length() == 12 else {
    fail("Invalid nonce length -- nonce must be 96 bits")
  }
  let key_ = FixedArray::make(8, 0U)
  for i in 0..<8 {
    key_[i] = (key[i * 4].to_uint() << 0) |
      (key[i * 4 + 1].to_uint() << 8) |
      (key[i * 4 + 2].to_uint() << 16) |
      (key[i * 4 + 3].to_uint() << 24)
  }
  let nonce_ = FixedArray::make(3, 0U)
  for i in 0..<3 {
    nonce_[i] = nonce[i * 4].to_uint() |
      (nonce[i * 4 + 1].to_uint() << 8) |
      (nonce[i * 4 + 2].to_uint() << 16) |
      (nonce[i * 4 + 3].to_uint() << 24)
  }
  {
    key: key_,
    nonce: nonce_,
    key_stream: FixedArray::make(64, b'\x00'),
    round,
    offset: 64,
    counter,
    exhausted: false,
  }
}

///|
/// Creates a ChaCha8 encryption context following the RFC 8439 standard.
/// - [key] must be 256-bit (32 bytes), in little-endian order.
/// - [nonce] must be a 96-bit (12 bytes) bytes, in little-endian order.
/// - [counter] is the counter value, defaulting to 0.
///
/// raise Error if the length of key or nonce is invalid.
pub fn[K : ByteSource, N : ByteSource] ChaCha::chacha8(
  key : K,
  nonce : N,
  counter? : UInt = 0,
) -> ChaCha raise Error {
  ChaCha::new(key, nonce, 4, counter)
}

///|
/// Creates a ChaCha12 encryption context following the RFC 8439 standard.
/// - [key] must be 256-bit (32 bytes), in little-endian order.
/// - [nonce] must be a 96-bit (12 bytes) bytes, in little-endian order.
/// - [counter] is the counter value, defaulting to 0.
///
/// raise Error if the length of key or nonce is invalid.
pub fn[K : ByteSource, N : ByteSource] ChaCha::chacha12(
  key : K,
  nonce : N,
  counter? : UInt = 0,
) -> ChaCha raise Error {
  ChaCha::new(key, nonce, 6, counter)
}

///|
/// Creates a ChaCha20 encryption context following the RFC 8439 standard.
/// - [key] must be 256-bit (32 bytes), in little-endian order.
/// - [nonce] must be a 96-bit (12 bytes) bytes, in little-endian order.
/// - [counter] is the counter value, defaulting to 0.
///
/// raise Error if the length of key or nonce is invalid.
pub fn[K : ByteSource, N : ByteSource] ChaCha::chacha20(
  key : K,
  nonce : N,
  counter? : UInt = 0,
) -> ChaCha raise Error {
  ChaCha::new(key, nonce, 10, counter)
}

///|
fn ChaCha::can_transform(self : ChaCha, data_length : Int) -> Bool {
  let buffered_bytes = if self.offset < 64 { 64 - self.offset } else { 0 }
  guard data_length > buffered_bytes else { true }
  guard !self.exhausted else { false }
  let remaining_bytes = data_length - buffered_bytes
  let required_blocks = (remaining_bytes - 1) / 64 + 1
  let available_blocks = 0x100000000UL - self.counter.to_uint64()
  required_blocks.to_uint64() <= available_blocks
}

///|
fn target_has_capacity(
  target : FixedArray[Byte],
  offset : Int,
  data_length : Int,
) -> Bool {
  guard offset >= 0 && offset <= target.length() else { false }
  data_length <= target.length() - offset
}

///|
fn ChaCha::generate_block(self : ChaCha) -> Unit {
  guard! !self.exhausted
  chachaBlock(self.key, self.counter, self.nonce, self.round, self.key_stream)
  if self.counter == 0xffffffff {
    self.exhausted = true
  } else {
    self.counter += 1
  }
}

///|
fn[D : ByteSource] ChaCha::transform_unchecked(
  self : Self,
  data : D,
  target : FixedArray[Byte],
  offset : Int,
  data_length : Int,
) -> Unit {
  let remaining_offset = if self.offset < 64 {
    let length = @cmp.minimum(data_length, 64 - self.offset)
    chacha_xor_into(
      data,
      0,
      target,
      offset,
      self.key_stream,
      self.offset,
      length,
    )
    self.offset += length
    length
  } else {
    0
  }
  for input_offset = remaining_offset; input_offset + 64 <= data_length; {
    self.generate_block()
    chacha_xor_into(
      data,
      input_offset,
      target,
      offset + input_offset,
      self.key_stream,
      0,
      64,
    )
    continue input_offset + 64
  } nobreak {
    if input_offset < data_length {
      self.generate_block()
      self.offset = data_length - input_offset
      chacha_xor_into(
        data,
        input_offset,
        target,
        offset + input_offset,
        self.key_stream,
        0,
        self.offset,
      )
    }
  }
}

///|
/// Transforms `data` with ChaCha and writes it into `target` at `offset`.
///
/// Aborts if `target` does not have room for the output or if the 32-bit block
/// counter would be exhausted. Use `transform_checked` when the caller needs
/// to handle these conditions.
pub fn[D : ByteSource] ChaCha::transform(
  self : Self,
  data : D,
  target : FixedArray[Byte],
  offset? : Int = 0,
) -> Unit {
  let data_length = data.length()
  guard target_has_capacity(target, offset, data_length) else {
    abort("invalid ChaCha target range")
  }
  guard self.can_transform(data_length) else {
    abort("ChaCha counter exhausted")
  }
  self.transform_unchecked(data, target, offset, data_length)
}

///|
/// Transforms `data` with ChaCha and writes it into `target` at `offset`.
///
/// Raises `ChaChaError::InvalidOutputRange` or
/// `ChaChaError::CounterExhausted` before writing any output if the target
/// range is invalid or the operation would exhaust the 32-bit block counter.
pub fn[D : ByteSource] ChaCha::transform_checked(
  self : Self,
  data : D,
  target : FixedArray[Byte],
  offset? : Int = 0,
) -> Unit raise ChaChaError {
  let data_length = data.length()
  guard target_has_capacity(target, offset, data_length) else {
    raise InvalidOutputRange
  }
  guard self.can_transform(data_length) else { raise CounterExhausted }
  self.transform_unchecked(data, target, offset, data_length)
}