defmodule Odinsea.Net.Cipher.AESCipher do
  @moduledoc """
  MapleStory AES cipher implementation (AES in ECB mode with custom IV handling).
  This is used for encrypting/decrypting packet data.

  Ported from: src/handling/netty/cipher/AESCipher.java
  """

  @block_size 1460

  # MapleStory AES key (32 bytes, expanded from the Java version)
  @aes_key <<
    0x13, 0x00, 0x00, 0x00,
    0x08, 0x00, 0x00, 0x00,
    0x06, 0x00, 0x00, 0x00,
    0xB4, 0x00, 0x00, 0x00,
    0x1B, 0x00, 0x00, 0x00,
    0x0F, 0x00, 0x00, 0x00,
    0x33, 0x00, 0x00, 0x00,
    0x52, 0x00, 0x00, 0x00
  >>

  @doc """
  Encrypts or decrypts packet data in place using AES-ECB with IV.

  ## Parameters
    - data: Binary data to encrypt/decrypt
    - iv: 4-byte IV binary

  ## Returns
    - Encrypted/decrypted binary data
  """
  @spec crypt(binary(), binary()) :: binary()
  def crypt(data, <<_::binary-size(4)>> = iv) when is_binary(data) do
    crypt_recursive(data, iv, 0, byte_size(data), @block_size - 4)
  end

  # Recursive encryption/decryption function
  @spec crypt_recursive(binary(), binary(), non_neg_integer(), non_neg_integer(), non_neg_integer()) :: binary()
  defp crypt_recursive(data, _iv, start, remaining, _length) when remaining <= 0 do
    # Return the portion of data we've processed
    binary_part(data, 0, start)
  end

  defp crypt_recursive(data, iv, start, remaining, length) do
    # Multiply the IV by 4
    seq_iv = multiply_bytes(iv, byte_size(iv), 4)

    # Adjust length if remaining is smaller
    actual_length = min(remaining, length)

    # Extract the portion of data to process
    data_bytes = :binary.bin_to_list(data)

    # Process the data chunk
    {new_data_bytes, _final_seq_iv} =
      process_chunk(data_bytes, seq_iv, start, start + actual_length, 0)

    # Convert back to binary
    new_data = :binary.list_to_bin(new_data_bytes)

    # Continue with next chunk
    new_start = start + actual_length
    new_remaining = remaining - actual_length
    new_length = @block_size

    crypt_recursive(new_data, iv, new_start, new_remaining, new_length)
  end

  # Process a single chunk of data
  @spec process_chunk(list(byte()), binary(), non_neg_integer(), non_neg_integer(), non_neg_integer()) ::
          {list(byte()), binary()}
  defp process_chunk(data_bytes, seq_iv, x, end_x, _offset) when x >= end_x do
    {data_bytes, seq_iv}
  end

  defp process_chunk(data_bytes, seq_iv, x, end_x, offset) do
    # Check if we need to re-encrypt the IV
    {new_seq_iv, new_offset} =
      if rem(offset, byte_size(seq_iv)) == 0 do
        # Encrypt the IV using AES
        encrypted_iv = aes_encrypt_block(seq_iv)
        {encrypted_iv, 0}
      else
        {seq_iv, offset}
      end

    # XOR the data byte with the IV byte
    seq_iv_bytes = :binary.bin_to_list(new_seq_iv)
    iv_index = rem(new_offset, length(seq_iv_bytes))
    iv_byte = Enum.at(seq_iv_bytes, iv_index)
    data_byte = Enum.at(data_bytes, x)
    xor_byte = Bitwise.bxor(data_byte, iv_byte)

    # Update the data
    updated_data = List.replace_at(data_bytes, x, xor_byte)

    # Continue processing
    process_chunk(updated_data, new_seq_iv, x + 1, end_x, new_offset + 1)
  end

  # Encrypt a single 16-byte block using AES-ECB
  @spec aes_encrypt_block(binary()) :: binary()
  defp aes_encrypt_block(block) do
    # Pad or truncate to 16 bytes for AES
    padded_block =
      case byte_size(block) do
        16 -> block
        size when size < 16 -> block <> :binary.copy(<<0>>, 16 - size)
        size when size > 16 -> binary_part(block, 0, 16)
      end

    # Perform AES encryption in ECB mode
    :crypto.crypto_one_time(:aes_128_ecb, @aes_key, padded_block, true)
  end

  # Multiply bytes - repeats the first `count` bytes of `input` `mul` times
  @spec multiply_bytes(binary(), non_neg_integer(), non_neg_integer()) :: binary()
  defp multiply_bytes(input, count, mul) do
    # Take first `count` bytes and repeat them `mul` times
    chunk = binary_part(input, 0, min(count, byte_size(input)))
    :binary.copy(chunk, mul)
  end
end