// 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.

// An MD5 message-digest algorithm implementation based on
// [RFC1321]       https://www.ietf.org/rfc/rfc1321.txt
// [Ron Rivest]    https://people.csail.mit.edu/rivest/Md5.c
// [md5-0.7.0]     https://docs.rs/md5/0.7.0/src/md5/lib.rs.html

///|
#alias(MD5Context, deprecated="Use `MD5` instead")
struct MD5 {
  reg : FixedArray[UInt] // state 'a' 'b' 'c' 'd'
  mut len : UInt64
  buf : FixedArray[Byte]
  mut buf_index : Int
}

///|
#inline
fn load_md5_words(block : FixedArray[Byte], input : FixedArray[UInt]) -> Unit {
  // Callers allocate the fixed MD5 block widths, so the unchecked accesses
  // below stay within block[0..64) and input[0..16).
  for word_offset in 0..<16 {
    let offset = word_offset * 4
    input.unsafe_set(
      word_offset,
      block.unsafe_get(offset).to_uint() |
      (block.unsafe_get(offset + 1).to_uint() << 8) |
      (block.unsafe_get(offset + 2).to_uint() << 16) |
      (block.unsafe_get(offset + 3).to_uint() << 24),
    )
  }
}

///|
pub impl CryptoHasher for MD5 with fn size(_self : MD5) -> Int {
  16
}

///|
pub impl CryptoHasher for MD5 with fn block_size(_self : MD5) -> Int {
  64
}

///|
pub impl CryptoHasher for MD5 with fn reset(self : MD5) -> Unit {
  self.reg[0] = 0x67452301
  self.reg[1] = 0xefcdab89
  self.reg[2] = 0x98badcfe
  self.reg[3] = 0x10325476
  self.len = 0
  self.buf.fill(0)
  self.buf_index = 0
}

///|
/// update the state of given context from new `data` 
pub fn[Data : ByteSource] MD5::update(self : MD5, data : Data) -> Unit {
  md5_update(self, data)
}

///|
pub impl CryptoHasher for MD5 with fn update(self : MD5, data : BytesView) -> Unit {
  md5_update(self, data)
}

///|
pub fn MD5::finalize(self : MD5) -> FixedArray[Byte] {
  self.md5_compute()
}

///|
/// an alias of `MD5Context::compute()`
pub impl CryptoHasher for MD5 with fn finalize_into(
  self : MD5,
  buffer : FixedArray[Byte],
  offset~ : Int,
) -> Unit {
  self.md5_compute().blit_to(buffer, len=16, src_offset=0, dst_offset=offset)
}

///|
/// Instantiate a MD5 context
pub fn MD5::new() -> MD5 {
  {
    reg: [0x67452301, 0xefcdab89, 0x98badcfe, 0x10325476],
    len: 0,
    buf: FixedArray::make(64, Byte::default()),
    buf_index: 0,
  }
}

///|
/// compute MD5 digest from given context
fn MD5::md5_compute(self : MD5) -> FixedArray[Byte] {
  let data = FixedArray::make(64, Byte::default())
  let mut cnt = self.buf_index
  let len = self.len + 8 * cnt.to_uint64()
  // Build the padded final block(s) on a copy so that `finalize` stays
  // idempotent and the context remains usable afterwards.
  self.buf.blit_to(data, len=cnt)
  let reg = self.reg.copy()
  data[cnt] = b'\x80'
  cnt += 1
  if cnt > 56 {
    let input = FixedArray::make(16, 0U)
    load_md5_words(data, input)
    md5_transform(reg, input)
    data.fill(0)
  }
  // little-endian 64-bit bit length at bytes 56..64
  for i in 0..<8 {
    data.unsafe_set(56 + i, (len >> (8 * i)).to_byte())
  }
  let input = FixedArray::make(16, 0U)
  load_md5_words(data, input)
  md5_transform(reg, input)
  let digest = FixedArray::make(16, Byte::default())
  // The MD5 state and digest have fixed four-word and 16-byte widths.
  for lane in 0..<4 {
    let word = reg.unsafe_get(lane)
    let offset = lane * 4
    digest.unsafe_set(offset, word.to_byte())
    digest.unsafe_set(offset + 1, (word >> 8).to_byte())
    digest.unsafe_set(offset + 2, (word >> 16).to_byte())
    digest.unsafe_set(offset + 3, (word >> 24).to_byte())
  }
  digest
}

// no macros, nor inline. basic md5 functions
// four auxiliary functions
//          F(X,Y,Z) = XY v not(X) Z
//          G(X,Y,Z) = XZ v Y not(Z)
//          H(X,Y,Z) = X xor Y xor Z
//          I(X,Y,Z) = Y xor (X v not(Z))

///|
fn MD5::f(x : UInt, y : UInt, z : UInt) -> UInt {
  (x & y) | (x.lnot() & z)
}

///|
fn MD5::g(x : UInt, y : UInt, z : UInt) -> UInt {
  ((x ^ y) & z) ^ y
}

///|
fn MD5::h(x : UInt, y : UInt, z : UInt) -> UInt {
  x ^ y ^ z
}

///|
fn MD5::i(x : UInt, y : UInt, z : UInt) -> UInt {
  y ^ (x | z.lnot())
}

///|
fn[Data : ByteSource] md5_update(ctx : MD5, data : Data) -> Unit {
  let input = FixedArray::make(16, 0U)
  let data_len = data.length()
  let mut offset = 0
  while offset < data_len {
    let chunk_length = @cmp.minimum(64 - ctx.buf_index, data_len - offset)
    data.blit_to(
      ctx.buf,
      len=chunk_length,
      src_offset=offset,
      dst_offset=ctx.buf_index,
    )
    ctx.buf_index += chunk_length
    offset += chunk_length
    if ctx.buf_index == 64 {
      ctx.len += 512UL
      ctx.buf_index = 0
      load_md5_words(ctx.buf, input)
      md5_transform(ctx.reg, input)
    }
  }
}

///|
#inline
fn md5_transform(state : FixedArray[UInt], input : FixedArray[UInt]) -> Unit {
  // load_md5_words always supplies exactly 16 words; every constant word index
  // used by the rounds below is therefore in bounds.
  let mut a = state[0]
  let mut b = state[1]
  let mut c = state[2]
  let mut d = state[3]

  // Round 1
  // s[ 0..15] := { 7, 12, 17, 22,  7, 12, 17, 22,  7, 12, 17, 22,  7, 12, 17, 22 }
  a = b +
    rotate_left_u(a + MD5::f(b, c, d) + input.unsafe_get(0) + 0xd76aa478, 7)
  d = a +
    rotate_left_u(d + MD5::f(a, b, c) + input.unsafe_get(1) + 0xe8c7b756, 12)
  c = d +
    rotate_left_u(c + MD5::f(d, a, b) + input.unsafe_get(2) + 0x242070db, 17)
  b = c +
    rotate_left_u(b + MD5::f(c, d, a) + input.unsafe_get(3) + 0xc1bdceee, 22)
  a = b +
    rotate_left_u(a + MD5::f(b, c, d) + input.unsafe_get(4) + 0xf57c0faf, 7)
  d = a +
    rotate_left_u(d + MD5::f(a, b, c) + input.unsafe_get(5) + 0x4787c62a, 12)
  c = d +
    rotate_left_u(c + MD5::f(d, a, b) + input.unsafe_get(6) + 0xa8304613, 17)
  b = c +
    rotate_left_u(b + MD5::f(c, d, a) + input.unsafe_get(7) + 0xfd469501, 22)
  a = b +
    rotate_left_u(a + MD5::f(b, c, d) + input.unsafe_get(8) + 0x698098d8, 7)
  d = a +
    rotate_left_u(d + MD5::f(a, b, c) + input.unsafe_get(9) + 0x8b44f7af, 12)
  c = d +
    rotate_left_u(c + MD5::f(d, a, b) + input.unsafe_get(10) + 0xffff5bb1, 17)
  b = c +
    rotate_left_u(b + MD5::f(c, d, a) + input.unsafe_get(11) + 0x895cd7be, 22)
  a = b +
    rotate_left_u(a + MD5::f(b, c, d) + input.unsafe_get(12) + 0x6b901122, 7)
  d = a +
    rotate_left_u(d + MD5::f(a, b, c) + input.unsafe_get(13) + 0xfd987193, 12)
  c = d +
    rotate_left_u(c + MD5::f(d, a, b) + input.unsafe_get(14) + 0xa679438e, 17)
  b = c +
    rotate_left_u(b + MD5::f(c, d, a) + input.unsafe_get(15) + 0x49b40821, 22)

  // Round 2
  // s[16..31] := { 5,  9, 14, 20,  5,  9, 14, 20,  5,  9, 14, 20,  5,  9, 14, 20 }
  a = b +
    rotate_left_u(a + MD5::g(b, c, d) + input.unsafe_get(1) + 0xf61e2562, 5)
  d = a +
    rotate_left_u(d + MD5::g(a, b, c) + input.unsafe_get(6) + 0xc040b340, 9)
  c = d +
    rotate_left_u(c + MD5::g(d, a, b) + input.unsafe_get(11) + 0x265e5a51, 14)
  b = c +
    rotate_left_u(b + MD5::g(c, d, a) + input.unsafe_get(0) + 0xe9b6c7aa, 20)
  a = b +
    rotate_left_u(a + MD5::g(b, c, d) + input.unsafe_get(5) + 0xd62f105d, 5)
  d = a +
    rotate_left_u(d + MD5::g(a, b, c) + input.unsafe_get(10) + 0x02441453, 9)
  c = d +
    rotate_left_u(c + MD5::g(d, a, b) + input.unsafe_get(15) + 0xd8a1e681, 14)
  b = c +
    rotate_left_u(b + MD5::g(c, d, a) + input.unsafe_get(4) + 0xe7d3fbc8, 20)
  a = b +
    rotate_left_u(a + MD5::g(b, c, d) + input.unsafe_get(9) + 0x21e1cde6, 5)
  d = a +
    rotate_left_u(d + MD5::g(a, b, c) + input.unsafe_get(14) + 0xc33707d6, 9)
  c = d +
    rotate_left_u(c + MD5::g(d, a, b) + input.unsafe_get(3) + 0xf4d50d87, 14)
  b = c +
    rotate_left_u(b + MD5::g(c, d, a) + input.unsafe_get(8) + 0x455a14ed, 20)
  a = b +
    rotate_left_u(a + MD5::g(b, c, d) + input.unsafe_get(13) + 0xa9e3e905, 5)
  d = a +
    rotate_left_u(d + MD5::g(a, b, c) + input.unsafe_get(2) + 0xfcefa3f8, 9)
  c = d +
    rotate_left_u(c + MD5::g(d, a, b) + input.unsafe_get(7) + 0x676f02d9, 14)
  b = c +
    rotate_left_u(b + MD5::g(c, d, a) + input.unsafe_get(12) + 0x8d2a4c8a, 20)

  // Round 3
  // s[32..47] := { 4, 11, 16, 23,  4, 11, 16, 23,  4, 11, 16, 23,  4, 11, 16, 23 }
  a = b +
    rotate_left_u(a + MD5::h(b, c, d) + input.unsafe_get(5) + 0xfffa3942, 4)
  d = a +
    rotate_left_u(d + MD5::h(a, b, c) + input.unsafe_get(8) + 0x8771f681, 11)
  c = d +
    rotate_left_u(c + MD5::h(d, a, b) + input.unsafe_get(11) + 0x6d9d6122, 16)
  b = c +
    rotate_left_u(b + MD5::h(c, d, a) + input.unsafe_get(14) + 0xfde5380c, 23)
  a = b +
    rotate_left_u(a + MD5::h(b, c, d) + input.unsafe_get(1) + 0xa4beea44, 4)
  d = a +
    rotate_left_u(d + MD5::h(a, b, c) + input.unsafe_get(4) + 0x4bdecfa9, 11)
  c = d +
    rotate_left_u(c + MD5::h(d, a, b) + input.unsafe_get(7) + 0xf6bb4b60, 16)
  b = c +
    rotate_left_u(b + MD5::h(c, d, a) + input.unsafe_get(10) + 0xbebfbc70, 23)
  a = b +
    rotate_left_u(a + MD5::h(b, c, d) + input.unsafe_get(13) + 0x289b7ec6, 4)
  d = a +
    rotate_left_u(d + MD5::h(a, b, c) + input.unsafe_get(0) + 0xeaa127fa, 11)
  c = d +
    rotate_left_u(c + MD5::h(d, a, b) + input.unsafe_get(3) + 0xd4ef3085, 16)
  b = c +
    rotate_left_u(b + MD5::h(c, d, a) + input.unsafe_get(6) + 0x04881d05, 23)
  a = b +
    rotate_left_u(a + MD5::h(b, c, d) + input.unsafe_get(9) + 0xd9d4d039, 4)
  d = a +
    rotate_left_u(d + MD5::h(a, b, c) + input.unsafe_get(12) + 0xe6db99e5, 11)
  c = d +
    rotate_left_u(c + MD5::h(d, a, b) + input.unsafe_get(15) + 0x1fa27cf8, 16)
  b = c +
    rotate_left_u(b + MD5::h(c, d, a) + input.unsafe_get(2) + 0xc4ac5665, 23)

  // Round 4
  // s[48..63] := { 6, 10, 15, 21,  6, 10, 15, 21,  6, 10, 15, 21,  6, 10, 15, 21 }
  a = b +
    rotate_left_u(a + MD5::i(b, c, d) + input.unsafe_get(0) + 0xf4292244, 6)
  d = a +
    rotate_left_u(d + MD5::i(a, b, c) + input.unsafe_get(7) + 0x432aff97, 10)
  c = d +
    rotate_left_u(c + MD5::i(d, a, b) + input.unsafe_get(14) + 0xab9423a7, 15)
  b = c +
    rotate_left_u(b + MD5::i(c, d, a) + input.unsafe_get(5) + 0xfc93a039, 21)
  a = b +
    rotate_left_u(a + MD5::i(b, c, d) + input.unsafe_get(12) + 0x655b59c3, 6)
  d = a +
    rotate_left_u(d + MD5::i(a, b, c) + input.unsafe_get(3) + 0x8f0ccc92, 10)
  c = d +
    rotate_left_u(c + MD5::i(d, a, b) + input.unsafe_get(10) + 0xffeff47d, 15)
  b = c +
    rotate_left_u(b + MD5::i(c, d, a) + input.unsafe_get(1) + 0x85845dd1, 21)
  a = b +
    rotate_left_u(a + MD5::i(b, c, d) + input.unsafe_get(8) + 0x6fa87e4f, 6)
  d = a +
    rotate_left_u(d + MD5::i(a, b, c) + input.unsafe_get(15) + 0xfe2ce6e0, 10)
  c = d +
    rotate_left_u(c + MD5::i(d, a, b) + input.unsafe_get(6) + 0xa3014314, 15)
  b = c +
    rotate_left_u(b + MD5::i(c, d, a) + input.unsafe_get(13) + 0x4e0811a1, 21)
  a = b +
    rotate_left_u(a + MD5::i(b, c, d) + input.unsafe_get(4) + 0xf7537e82, 6)
  d = a +
    rotate_left_u(d + MD5::i(a, b, c) + input.unsafe_get(11) + 0xbd3af235, 10)
  c = d +
    rotate_left_u(c + MD5::i(d, a, b) + input.unsafe_get(2) + 0x2ad7d2bb, 15)
  b = c +
    rotate_left_u(b + MD5::i(c, d, a) + input.unsafe_get(9) + 0xeb86d391, 21)
  state[0] += a
  state[1] += b
  state[2] += c
  state[3] += d
}

///|
/// Compute the MD5 digest of some `data` based on [RFC1321](https://www.ietf.org/rfc/rfc1321.txt).
/// - Note that MD5 is considered _cryptographically broken_.
/// Unless mandated, more secure alternatives should be preferred.
pub fn[Data : ByteSource] md5(data : Data) -> FixedArray[Byte] {
  let ctx = MD5::new()
  md5_update(ctx, data)
  ctx.md5_compute()
}

///|
test {
  let ctx = MD5::new()
  md5_update(ctx, b"\x61".to_fixedarray())
  md5_update(ctx, b"\x62".to_fixedarray())
  md5_update(ctx, b"\x63".to_fixedarray())
  let res1 = bytes_to_hex_string(ctx.md5_compute())
  let ctx = MD5::new()
  md5_update(ctx, b"\x61\x62\x63".to_fixedarray())
  let res2 = bytes_to_hex_string(ctx.md5_compute())
  assert_eq(res1, res2)
}

///|
#deprecated
pub extend MD5 with CryptoHasher::{reset, finalize_into, size, block_size}