reticulum

defmodule RNS.Packet do
  @moduledoc """
  Everything related to Reticulum packets.
  NOTE: Instead of decoding packets to a struct, you just fetch the part you want out of the binary data.

  ## Fields of a packet
  Header field:
  - Interface Access Controls(IFAC) flag: open or authenticated, authenticated means there will be an additional ifac field.
  - Header type: Basically if it's set to true there are 2 destination fields, second one is transport id.
  - Context flag: depends on context...
  - Propagation type: Can be transport or broadcast.
  - Destination type: Can be single, group, plain or link.
  - Packet type: can be data, announce, link request or proof.
  - Hops: the number of hops the packet has gone through(it is increased when received).

  (only if the IFAC flag is set) IFAC field
  - IFAC: Not implemented yet, it encrypts the packet with AES.

  Destination field
  - Destination: the destination of the packet.
  - (only if header type is of type 2(set)) Transport ID: idk much.

  Context field
  - Context: a byte, and contains more information related to the context, some examples are 0x00 for none(default) or 0x03 for a resource request.

  Data field
  Contains the actual data about the packet.
  """
  @type t() :: tuple()
  @type raw() :: <<_::152, _::_*8>>
  # Packet Fields
  @type header() :: <<_::16>>
  @type context() :: 0x00..0xFF
  @type data() :: binary()

  # Header Fields
  @type ifac_flag() :: :open | :authenticated
  @type header_type() :: boolean()
  @type context_flag() :: boolean()
  @type propagation_type() :: :broadcast | :transport
  @type packet_type() :: :data | :announce | :link_request | :proof
  @type hops() :: 0x00..0xFF

  # Packet Metadata
  @type hash() :: RNS.Crypto.hash()
  @type truncated_hash() :: RNS.Crypto.truncated_hash()

  # NOTE: Implement MTU and MDU!
  # The maximum size of a packet.
  @mtu 500

  # The maximum size of the data field.
  # The mtu minus: header(2) + content(1) + destinations(max 32) + the minimum size of the ifac(1)
  @mdu 472

  require Logger

  require Record
  Record.defrecord(:packet,
    # Packet Fields
    destination_hash: nil,
    transport_id: nil,
    context: 0x00,
    data: nil,
    encrypted_data: nil,
    # Header Fields
    header_type: false,
    context_flag: false,
    propagation_type: :transport,
    destination_type: :single,
    packet_type: :data,
    hops: 0,
    # Packet Metadata
    hash: nil,
    raw: nil
  )
  # Packet Fields

  @doc "Get the destination hash of a packet record."
  @spec destination_hash(t()) :: RNS.Destination.hash()
  def destination_hash(packet), do: packet(packet, :destination_hash)
  
  @doc "Get the transport id of a packet record."
  @spec transport_id(t()) :: RNS.Destination.hash() | nil
  def transport_id(packet), do: packet(packet, :transport_id)

  @doc "Get the context of a packet record."
  @spec context(t()) :: context()
  def context(packet), do: packet(packet, :context)

  @doc "Get the data of a packet record."
  @spec data(t()) :: data() | nil
  def data(packet), do: packet(packet, :data)

  @doc "Get the encrypted data of a packet record."
  @spec encrypted_data(t()) :: data() | nil
  def encrypted_data(packet), do: packet(packet, :encrypted_data)

  # Header Fields

  @doc "Get the header type of a packet record."
  @spec header_type(t()) :: header_type()
  def header_type(packet), do: packet(packet, :header_type)

  @doc "Get the context flag of a packet record."
  @spec context_flag(t()) :: context_flag()
  def context_flag(packet), do: packet(packet, :context_flag)

  @doc "Get the propagation type of a packet record."
  @spec propagation_type(t()) :: propagation_type()
  def propagation_type(packet), do: packet(packet, :propagation_type)

  @doc "Get the destination type of a packet record."
  @spec destination_type(t()) :: RNS.Destination.type()
  def destination_type(packet), do: packet(packet, :destination_type)

  @doc "Get the packet type of a packet record."
  @spec packet_type(t()) :: packet_type()
  def packet_type(packet), do: packet(packet, :packet_type)

  @doc "Get the hops of a packet record."
  @spec hops(t()) :: hops()
  def hops(packet), do: packet(packet, :hops)
  
  # Packet Metadata

  @doc "Get the hash of a packet record."
  @spec hash(t()) :: hash() | nil
  def hash(packet), do: packet(packet, :hash)

  @doc "Get the truncated hash of a packet record."
  @spec truncated_hash(t()) :: truncated_hash() | nil
  def truncated_hash(packet) do
    packet(packet, :hash)
    |> RNS.Crypto.truncate(16)
  end

  @doc "Get the raw packet of a packet record."
  @spec raw(t()) :: raw()
  def raw(packet), do: packet(packet, :raw)


  # NOTE: the header type is done automatically, and IFAC should be handled while encoding/decoding the packet.
  @doc """
  Creates a new packet.
  The packet's destination and data is required.
  Options related to the packet:
  - transport_id: A 16 byte binary. If it is omitted, 
  - context: An integer between 0x00 and 0xFF, 0x00 by default. For more information, see `RNS.PacketContext`.
  Options related to the header:
  - context_flag: Depends on the context... Announces use it to say if a ratche is announced or not.
  - propagation_type: How the message will be propagated, :transport means the packet will be routed, :broadcast means it will be sent to all interfaces.
  - destination_type: The type of the destination, :single means the data will be encrypted and the message is only to one node, :group means it is encrypted but sent to multiple nodes, :plain means the content is not encrypted, for example for an announce, and :link means it's for a link.
  - packet_type: The type of the packet, default is :data, just a normal data packet. Can also be :announce, :link_request or :proof.
  - hops: the number of hops the packet has been routed by, it is incremented when received.
  """
  @spec new(
          # Packet Related
          data(),
          RNS.Destination.hash(),
          [
            transport_id: RNS.Destination.hash(),
            context: context(),
            # Header Related
            context_flag: context_flag(),
            propagation_type: propagation_type(),
            destination_type: RNS.Destination.type(),
            packet_type: packet_type(),
            hops: hops()
          ]
        ) :: t()
  def new(data, destination_hash, opts \\ []) do
    transport_id = Keyword.get(opts, :transport_id, nil)
    context = Keyword.get(opts, :context, 0x00)
    
    header_type = (if transport_id == nil, do: false, else: true)
    context_flag = Keyword.get(opts, :context_flag, false)
    propagation_type = Keyword.get(opts, :propagation_type, :transport)
    destination_type = Keyword.get(opts, :destination_type, :single)
    packet_type = Keyword.get(opts, :packet_type, :data)
    hops = Keyword.get(opts, :hops, 0)

    packet(
      # Packet Fields
      destination_hash: destination_hash,
      transport_id: transport_id,
      context: context,
      data: data,
      # Header Fields
      header_type: header_type,
      context_flag: context_flag,
      propagation_type: propagation_type,
      destination_type: destination_type,
      packet_type: packet_type,
      hops: hops
    )
  end

  @doc "Generates the raw packet from a packet record(returns the updated packet record)."
  @spec generate_raw(t()) :: t()
  def generate_raw(packet) do
    transport_id = transport_id(packet)
  
    raw_destination_type =
      case destination_type(packet) do
        :single -> 0b00
        :group ->  0b01
        :plain ->  0b10
        :link ->   0b11 
      end

    raw_packet_type =
      case packet_type(packet) do
        :data -> 0b00
        :announce -> 0b01
        :link_request -> 0b10
        :proof -> 0b11
      end
    
    raw = 
      <<
        0::1,
        (if header_type(packet), do: 1, else: 0)::1,
        (if context_flag(packet), do: 1, else: 0)::1,
        (if propagation_type(packet) == :transport, do: 1, else: 0)::1,
        raw_destination_type::2,
        raw_packet_type::2,
        hops(packet)::8,
        destination_hash(packet)::binary-size(16),
        (if transport_id == nil, do: <<>>, else: transport_id)::binary,
        context(packet)::8,
        encrypted_data(packet)::binary
      >>

    hashless_packet = packet(packet, raw: raw)
    packet(hashless_packet, hash: generate_hash(hashless_packet))
  end

  @doc "Creates a packet record from a raw packet."
  @spec from_raw(raw()) :: t()
  def from_raw(raw) do
    <<
      _raw_ifac_flag::1,
      raw_header_type::1,
      raw_context_flag::1,
      raw_propagation_type::1,
      raw_destination_type::2,
      raw_packet_type::2,
      hops::8,
      not_header::binary
    >> = raw

    header_type = (if raw_header_type == 1, do: true, else: false)
    context_flag = (if raw_context_flag == 1, do: true, else: false)
    propagation_type = (if raw_propagation_type == 1, do: :transport, else: :broadcast)
    destination_type =
      case raw_destination_type do
        0b00 -> :single
        0b01 -> :group
        0b10 -> :plain
        0b11 -> :link
      end
    packet_type =
      case raw_packet_type do
        0b00 -> :data
        0b01 -> :announce
        0b10 -> :link_request
        0b11 -> :proof
      end

    {destination_hash, transport_id, context, encrypted_data} =
      if header_type do
        <<
          destination_hash::binary-size(16),
          transport_id::binary-size(16),
          context::8,
          encrypted_data::binary
        >> = not_header
        {destination_hash, transport_id, context, encrypted_data}
      else
        <<
          destination_hash::binary-size(16),
          context::8,
          encrypted_data::binary
        >> = not_header
        {destination_hash, nil, context, encrypted_data}
      end

    hashless_packet = packet(
        # Packet Fields
        destination_hash: destination_hash,
        transport_id: transport_id,
        context: context,
        encrypted_data: encrypted_data,
        # Header Fields
        header_type: header_type,
        context_flag: context_flag,
        propagation_type: propagation_type,
        destination_type: destination_type,
        packet_type: packet_type,
        hops: hops,
        # Packet Metadata
        raw: raw
      )

    packet(hashless_packet, hash: generate_hash(hashless_packet))
  end

  
  @doc "Returns the MTU of a Reticulum packet(max size)."
  @spec mtu() :: pos_integer()
  def mtu(), do: @mtu

  @doc "Returns the MDU of a Reticulum packet(max size of the data field)."
  @spec mdu() :: pos_integer()
  def mdu(), do: @mdu

  @doc "Get the hashable part of a packet."
  @spec hashable_part(t()) :: binary()
  def hashable_part(packet) do
    <<first_byte::8, _::8, remains::binary>> = raw(packet)
    hashable_first_byte = Bitwise.band(first_byte, 0b00001111)

    case header_type(packet) do
      false ->
        <<hashable_first_byte::8, remains::binary>>

      true ->
        <<_::binary-size(16), remaining_remains::binary>> = remains
        <<hashable_first_byte::8, remaining_remains::binary>>
    end
  end

  @doc "Returns the packet hash in it's 32 bytes form."
  @spec generate_hash(t()) :: hash()
  def generate_hash(packet) do
    packet
    |> hashable_part()
    |> RNS.Crypto.sha256()
  end

  @spec decrypt(t(), RNS.Identity.t()) :: {:ok, t()} | {:error, atom()}
  def decrypt(packet, identity) do
    case RNS.Identity.decrypt(encrypted_data(packet), identity) do
      {:ok, data} ->
        {:ok, packet(packet, data: data)}
      {:error, reason} ->
        {:error, reason}
    end
  end

  @spec encrypt(t(), RNS.Identity.t(), RNS.Identity.ratchet() | nil) :: t()
  def encrypt(packet, identity, ratchet) do
    encrypted_data = RNS.Identity.encrypt(data(packet), identity, ratchet)
    packet(packet, encrypted_data: encrypted_data)
  end

  @doc """
  Encrypts(when needed) and sends a packet, then waits for the packet proof.
  This should not be directly ran by a process which could receive messages, because it uses `receive`; please run this as a task instead!
  """
  @spec send(t(), RNS.Destination.t()) ::
          :ok
          | {:error,
             :no_receipt | :identity_required | :unsupported_destination_type | :unreachable}
  def send(packet, destination) do
    destination_identity = RNS.Destination.identity(destination)

    if destination_identity == nil do
      if destination_type(packet) == :plain do
        packet
        |> generate_raw()
        |> send_encrypted(destination)
      else
        {:error, :identity_required}
      end
    else
      case destination_type(packet) do
        :plain ->
          packet
          |> generate_raw()
          |> send_encrypted(destination)

        :single ->
          ratchet =
            case RNS.DestinationStore.fetch_recent_ratchet(RNS.Destination.hash(destination)) do
              {:ok, ratchet} ->
                ratchet

              {:error, :not_found} ->
                nil
            end

          packet
          |> encrypt(destination_identity, ratchet)
          |> generate_raw()
          |> send_encrypted(destination)

        :group ->
          # TODO: Add support for group destinations.
          {:error, :unsupported_destination_type}

        :link ->
          # TODO: Add support for links.
          {:error, :unsupported_destination_type}
      end
    end
  end

  @spec send_encrypted(t(), RNS.Destination.t()) :: :ok | {:error, :timeout | :unreachable}
  defp send_encrypted(packet, destination) do
    # Add the destination to listen for the proof to the store.
    hash = hash(packet)
    truncated_hash = truncated_hash(packet)
    RNS.Destination.new(hash: truncated_hash, interface: self(), store: true)

    interface = RNS.Destination.interface(destination)

    case RNS.Interface.send(packet, interface, self()) do
      :ok ->
        status = wait_for_proof(hash, destination)
        # Unregister the destination.
        RNS.DestinationStore.delete(truncated_hash)
        status

      {:error, :dead_interface} ->
        # Unregister the destination.
        RNS.DestinationStore.delete(truncated_hash)
        {:error, :unreachable}

      results ->
        # If any result is ok, wait for a response, else return the :unreachable error.
        if Enum.any?(results, fn result -> if result == :ok, do: true, else: false end) do
          status = wait_for_proof(hash, destination)
          # Unregister the destination.
          RNS.DestinationStore.delete(truncated_hash)
          status
        else
          # Unregister the destination.
          RNS.DestinationStore.delete(truncated_hash)
          {:error, :unreachable}
        end
    end
  end

  # Waits for a proof, with a timeout of 30 seconds.
  @spec wait_for_proof(RNS.Packet.hash(), RNS.Destination.t(), NaiveDateTime.t()) ::
          :ok | {:error, :timeout}
  defp wait_for_proof(
         packet_hash,
         destination,
         timeout \\ NaiveDateTime.add(NaiveDateTime.utc_now(), 30, :second)
       ) do
    destination_identity = RNS.Destination.identity(destination)

    receive do
      {:inbound, packet, source} ->
        # TODO: check RNS.Destination.interface(destination) == source  but check if id, pid, record.
        if RNS.Packet.DataPacketProof.validate?(packet, packet_hash, destination_identity) do
          :ok
        else
          Logger.warning(
            "Received invalid data packet proof for data packet #{Base.encode16(packet_hash)}!"
          )
        end

      _ ->
        wait_for_proof(packet_hash, destination_identity, timeout)
    after
      NaiveDateTime.diff(timeout, NaiveDateTime.utc_now(), :millisecond) ->
        {:error, :timeout}
    end
  end
end