eyeon

defmodule Eyeon.Binary do
  @moduledoc """
  Shared primitives for Ion binary encoding: VarUInt, VarInt, UInt, Int.

  Ion binary uses big-endian byte order. VarUInt/VarInt use 7 bits per byte
  with the high bit as a stop bit (1 = last byte).
  """

  import Bitwise

  # --- Decoding ---

  @spec read_varuint(nonempty_binary()) :: {non_neg_integer(), binary()}
  def read_varuint(data), do: do_read_varuint(data, 0)

  defp do_read_varuint(<<1::1, value::7, rest::binary>>, acc) do
    {bsl(acc, 7) + value, rest}
  end

  defp do_read_varuint(<<0::1, value::7, rest::binary>>, acc) do
    do_read_varuint(rest, bsl(acc, 7) + value)
  end

  @spec read_varint(nonempty_binary()) :: {integer(), binary()}
  def read_varint(<<1::1, sign::1, value::6, rest::binary>>) do
    result = value
    if sign == 1, do: {-result, rest}, else: {result, rest}
  end

  def read_varint(<<0::1, sign::1, value::6, rest::binary>>) do
    do_read_varint(rest, value, sign)
  end

  defp do_read_varint(<<1::1, value::7, rest::binary>>, acc, sign) do
    result = bsl(acc, 7) + value
    if sign == 1, do: {-result, rest}, else: {result, rest}
  end

  defp do_read_varint(<<0::1, value::7, rest::binary>>, acc, sign) do
    do_read_varint(rest, bsl(acc, 7) + value, sign)
  end

  @spec read_uint(binary(), non_neg_integer()) :: {non_neg_integer(), binary()}
  def read_uint(data, 0), do: {0, data}

  def read_uint(data, len) do
    <<bytes::binary-size(len), rest::binary>> = data
    {:binary.decode_unsigned(bytes, :big), rest}
  end

  @spec read_int(binary(), non_neg_integer()) :: {integer(), binary()}
  def read_int(data, 0), do: {0, data}

  def read_int(data, len) do
    <<first_byte, tail::binary>> = data
    remaining_len = len - 1
    <<remaining::binary-size(remaining_len), rest::binary>> = tail

    sign = bsr(first_byte, 7)
    magnitude_first = band(first_byte, 0x7F)

    magnitude =
      if remaining_len == 0 do
        magnitude_first
      else
        mag_bytes = <<magnitude_first, remaining::binary>>
        :binary.decode_unsigned(mag_bytes, :big)
      end

    result = if sign == 1 and magnitude != 0, do: -magnitude, else: magnitude
    {result, rest}
  end

  # --- Encoding ---

  @spec encode_varuint(non_neg_integer()) :: binary()
  def encode_varuint(value) when value >= 0 do
    # Build bytes from least significant to most significant
    # Last transmitted byte (least significant) has stop bit (0x80)
    chunks = split_varuint_chunks(value)

    case chunks do
      [single] ->
        <<bor(single, 0x80)>>

      list ->
        {init, [last]} = Enum.split(list, -1)
        bytes = Enum.map(init, &band(&1, 0x7F)) ++ [bor(last, 0x80)]
        :erlang.list_to_binary(bytes)
    end
  end

  defp split_varuint_chunks(value) do
    do_split_varuint(value, [])
  end

  defp do_split_varuint(value, []) do
    chunk = band(value, 0x7F)
    remaining = bsr(value, 7)

    if remaining == 0 do
      [chunk]
    else
      do_split_varuint(remaining, [chunk])
    end
  end

  defp do_split_varuint(value, acc) do
    chunk = band(value, 0x7F)
    remaining = bsr(value, 7)

    if remaining == 0 do
      [chunk | acc]
    else
      do_split_varuint(remaining, [chunk | acc])
    end
  end

  @spec encode_varint(integer()) :: binary()
  def encode_varint(value) do
    {sign, magnitude} = if value < 0, do: {1, -value}, else: {0, value}

    # Determine how many bytes we need.
    # First byte has 6 data bits, subsequent have 7.
    # Capacity: 1 byte = 6 bits (max 63)
    #           2 bytes = 6 + 7 = 13 bits
    #           n bytes = 6 + 7*(n-1) bits
    #
    # Byte layout (MSB first in transmission):
    #   Byte 0: [stop?] [sign] [6 data bits - MSB of magnitude]
    #   Byte 1..n-2: [0] [7 data bits]
    #   Byte n-1: [1] [7 data bits - LSB of magnitude]
    #
    # For single byte: [1] [sign] [6 data bits]

    if magnitude <= 0x3F do
      <<1::1, sign::1, magnitude::6>>
    else
      # Split into 7-bit groups from LSB, then take up to 6 for the MSB byte
      groups = split_7bit_lsb_first(magnitude)
      # groups = [lsb_7bits, ..., msb_remaining_bits]

      # The last group (MSB) must fit in 6 bits for the first byte.
      # If it doesn't, we need to split it.
      {lsb_groups, msb} = split_for_varint_first_byte(groups)

      # Build bytes: first byte has sign + msb (6 bits), last byte has stop bit
      first_byte = bsl(sign, 6) ||| band(msb, 0x3F)
      last_byte = bor(0x80, band(hd(lsb_groups), 0x7F))

      case tl(lsb_groups) do
        [] ->
          <<first_byte, last_byte>>

        middle_groups ->
          middle_bytes = Enum.map(Enum.reverse(middle_groups), &band(&1, 0x7F))
          :erlang.list_to_binary([first_byte | middle_bytes] ++ [last_byte])
      end
    end
  end

  defp split_7bit_lsb_first(0), do: [0]

  defp split_7bit_lsb_first(value) do
    do_split_7bit_lsb(value, [])
  end

  defp do_split_7bit_lsb(0, acc), do: acc

  defp do_split_7bit_lsb(value, acc) do
    do_split_7bit_lsb(bsr(value, 7), [band(value, 0x7F) | acc])
  end

  # Returns {lsb_groups_in_msb_first_order, msb_value_fitting_in_6_bits}
  defp split_for_varint_first_byte(groups) do
    # groups is in [lsb, ..., msb] order (MSB last from our split)
    # Actually from do_split_7bit_lsb, groups accumulates with prepend,
    # so it's [msb, ..., lsb] order
    [msb | rest] = groups

    if msb <= 0x3F do
      {rest, msb}
    else
      # MSB group is > 6 bits, split it
      low = band(msb, 0x7F)
      high = bsr(msb, 7)
      # Replace with two groups
      split_for_varint_first_byte([high, low | rest])
    end
  end

  @spec encode_uint(non_neg_integer()) :: binary()
  def encode_uint(0), do: <<>>

  def encode_uint(value) when value > 0 do
    :binary.encode_unsigned(value, :big)
  end

  @spec encode_int(integer()) :: binary()
  def encode_int(0), do: <<>>

  def encode_int(value) do
    {sign_bit, magnitude} = if value < 0, do: {1, -value}, else: {0, value}
    mag_bytes = :binary.encode_unsigned(magnitude, :big)

    <<first_byte, rest::binary>> = mag_bytes

    if band(first_byte, 0x80) != 0 do
      # High bit already used, need extra byte for sign
      <<bsl(sign_bit, 7), first_byte, rest::binary>>
    else
      <<bor(bsl(sign_bit, 7), first_byte), rest::binary>>
    end
  end
end