nostr_lib

defmodule Nostr.Message do
  @moduledoc """
  Nostr message is represented as tuple in Elixir. This module provides functions to generate
  messages, serialize or parse them.
  """

  require Logger

  @type t() ::
          {:event, Nostr.Event.t()}
          | {:event, binary(), Nostr.Event.t()}
          | {:req, binary(), [Nostr.Filter.t()]}
          | {:close, binary()}
          | {:neg_open, binary(), Nostr.Filter.t(), binary()}
          | {:neg_msg, binary(), binary()}
          | {:neg_close, binary()}
          | {:neg_err, binary(), binary()}
          | {:eose, binary()}
          | {:notice, String.t()}
          | {:ok, binary(), boolean(), String.t()}
          | {:auth, Nostr.Event.t() | binary()}
          | {:count, String.t(), Nostr.Filter.t() | [Nostr.Filter.t()] | count_payload()}
          | {:closed, String.t(), String.t()}

  @type count_payload() :: %{
          required(:count) => integer(),
          optional(:approximate) => boolean(),
          optional(:hll) => binary()
        }

  @type parse_reason() ::
          :invalid_message_format | :unsupported_json_escape | :unsupported_json_literals

  @doc """
  Generate post new event message
  """
  @doc sender: :client
  @spec create_event(Nostr.Event.t() | %{event: Nostr.Event.t()}) :: {:event, Nostr.Event.t()}
  def create_event(%{event: %Nostr.Event{} = event}), do: {:event, event}
  def create_event(%Nostr.Event{} = event), do: {:event, event}

  @doc """
  Generate request message
  """
  @doc sender: :client
  @spec request(Nostr.Filter.t() | [Nostr.Filter.t()], binary()) ::
          {:req, binary(), [Nostr.Filter.t()]}
  def request(%Nostr.Filter{} = filter, sub_id), do: {:req, sub_id, [filter]}
  def request(filters, sub_id) when is_list(filters), do: {:req, sub_id, filters}

  @doc """
  Generate close message
  """
  @doc sender: :client
  @spec close(binary()) :: {:close, binary()}
  def close(sub_id), do: {:close, sub_id}

  @doc """
  Generate negentropy open message (NIP-77)
  """
  @doc sender: :client
  @spec neg_open(binary(), Nostr.Filter.t(), binary()) ::
          {:neg_open, binary(), Nostr.Filter.t(), binary()}
  def neg_open(sub_id, %Nostr.Filter{} = filter, initial_message),
    do: {:neg_open, sub_id, filter, initial_message}

  @doc """
  Generate negentropy message frame (NIP-77)
  """
  @doc sender: [:client, :relay]
  @spec neg_msg(binary(), binary()) :: {:neg_msg, binary(), binary()}
  def neg_msg(sub_id, message), do: {:neg_msg, sub_id, message}

  @doc """
  Generate negentropy close message (NIP-77)
  """
  @doc sender: :client
  @spec neg_close(binary()) :: {:neg_close, binary()}
  def neg_close(sub_id), do: {:neg_close, sub_id}

  @doc """
  Generate negentropy error message (NIP-77)
  """
  @doc sender: :relay
  @spec neg_err(binary(), binary()) :: {:neg_err, binary(), binary()}
  def neg_err(sub_id, reason), do: {:neg_err, sub_id, reason}

  @doc """
  Generate count message (NIP-45).

  Can be used to request counts from relay (with filters) or respond with counts (with integer).

  Parsed COUNT responses preserve optional `approximate` and `hll` keys when present.
  """
  @spec count(pos_integer() | Nostr.Filter.t() | [Nostr.Filter.t()], binary()) ::
          {:count, binary(), count_payload() | Nostr.Filter.t() | [Nostr.Filter.t()]}
  def count(count, sub_id) when is_integer(count), do: {:count, sub_id, %{count: count}}
  def count(%Nostr.Filter{} = filter, sub_id), do: {:count, sub_id, filter}
  def count(filters, sub_id) when is_list(filters), do: {:count, sub_id, filters}

  @doc """
  Generate event message
  """
  @doc sender: :relay
  @spec event(Nostr.Event.t() | %{event: Nostr.Event.t()}, binary()) ::
          {:event, binary(), Nostr.Event.t()}
  def event(%{event: %Nostr.Event{} = event}, sub_id), do: {:event, sub_id, event}
  def event(%Nostr.Event{} = event, sub_id), do: {:event, sub_id, event}

  @doc """
  Generate notice message
  """
  @doc sender: :relay
  @spec notice(String.t()) :: {:notice, String.t()}
  def notice(message), do: {:notice, message}

  @doc """
  Generate eose message
  """
  @doc sender: :relay
  @spec eose(binary()) :: {:eose, binary()}
  def eose(sub_id), do: {:eose, sub_id}

  @doc """
  Generate OK message
  """
  @doc sender: :relay
  @spec ok(binary(), boolean(), String.t()) :: {:ok, binary(), boolean(), String.t()}
  def ok(event_id, success?, message), do: {:ok, event_id, success?, message}

  @doc """
  Generate CLOSED message (NIP-01)
  """
  @doc sender: :relay
  @spec closed(String.t(), String.t()) :: {:closed, String.t(), String.t()}
  def closed(sub_id, message), do: {:closed, sub_id, message}

  @doc """
  Generate AUTH message (NIP-42).

  Can be used by relay (with challenge string) or by client (with signed event).
  """
  @spec auth(Nostr.Event.t() | %{event: Nostr.Event.t()} | binary()) ::
          {:auth, Nostr.Event.t() | binary()}
  def auth(%{event: %Nostr.Event{} = event}), do: {:auth, event}
  def auth(%Nostr.Event{} = event), do: {:auth, event}
  def auth(challenge), do: {:auth, challenge}

  @doc """
  Serialize Elixir tuple message to on-the-wire binary
  """
  @spec serialize(tuple()) :: binary()
  def serialize(message) when is_tuple(message) do
    message
    |> Tuple.to_list()
    |> List.flatten()
    |> then(fn [name | rest] ->
      name_str =
        name
        |> Atom.to_string()
        |> String.upcase()
        |> String.replace("_", "-")

      [name_str | rest]
    end)
    |> JSON.encode!()
  end

  @doc """
  Parse binary message to Elixir tuple, if message contains event it will be returned as general
  `Nostr.Event.t()` struct.

  Returns `:error` for invalid payloads.
  """
  @spec parse(msg :: String.t()) :: t()
  def parse(msg) when is_binary(msg) do
    case parse_with_reason(msg) do
      {:ok, parsed_message} -> parsed_message
      {:error, _reason} -> :error
    end
  end

  @doc """
  Parse binary message to Elixir tuple and return structured parse errors.

  The protocol-level parser in the relay uses this function to distinguish malformed JSON
  escape/literal cases from other parse failures.
  """
  @spec parse_with_reason(String.t()) :: {:ok, t()} | {:error, parse_reason()}
  def parse_with_reason(msg) when is_binary(msg) do
    with :ok <- validate_json_payload(msg),
         {:ok, decoded} <- JSON.decode(msg) do
      case do_parse(decoded, :general) do
        :error -> {:error, :invalid_message_format}
        parsed -> {:ok, parsed}
      end
    else
      {:error, reason} when reason in [:unsupported_json_escape, :unsupported_json_literals] ->
        {:error, reason}

      {:error, _decode_error} ->
        {:error, :invalid_message_format}
    end
  end

  defp validate_json_payload(msg) when is_binary(msg) do
    scan_json_payload(msg, false, false)
  end

  defp scan_json_payload(<<>>, false, false), do: :ok
  defp scan_json_payload(<<>>, true, false), do: {:error, :invalid_message_format}
  defp scan_json_payload(<<>>, _in_string, true), do: {:error, :unsupported_json_escape}

  defp scan_json_payload(<<?\", rest::binary>>, false, false),
    do: scan_json_payload(rest, true, false)

  defp scan_json_payload(<<?\\, rest::binary>>, true, false),
    do: scan_json_payload(rest, true, true)

  defp scan_json_payload(<<?\", rest::binary>>, true, false),
    do: scan_json_payload(rest, false, false)

  defp scan_json_payload(<<byte::8, _rest::binary>>, true, false) when byte < 32,
    do: {:error, :unsupported_json_literals}

  defp scan_json_payload(<<_byte::8, rest::binary>>, true, false),
    do: scan_json_payload(rest, true, false)

  defp scan_json_payload(<<_byte::8, rest::binary>>, false, false),
    do: scan_json_payload(rest, false, false)

  defp scan_json_payload(<<?", rest::binary>>, true, true),
    do: scan_json_payload(rest, true, false)

  defp scan_json_payload(<<?\\, ?", _rest::binary>>, true, true),
    do: {:error, :unsupported_json_escape}

  defp scan_json_payload(<<byte::8, rest::binary>>, true, true)
       when byte in [34, 92, ?b, ?f, ?n, ?r, ?t],
       do: scan_json_payload(rest, true, false)

  defp scan_json_payload(<<?u, a::8, b::8, c::8, d::8, rest::binary>>, true, true) do
    if valid_hex_byte?(a) and valid_hex_byte?(b) and valid_hex_byte?(c) and valid_hex_byte?(d) do
      if json_unicode_escape_codepoint(a, b, c, d) < 32 do
        {:error, :unsupported_json_escape}
      else
        scan_json_payload(rest, true, false)
      end
    else
      {:error, :unsupported_json_escape}
    end
  end

  defp scan_json_payload(<<?\", _rest::binary>>, true, true),
    do: {:error, :unsupported_json_escape}

  defp scan_json_payload(<<?\\, _rest::binary>>, true, true),
    do: {:error, :unsupported_json_escape}

  defp scan_json_payload(<<?\/, _rest::binary>>, true, true),
    do: {:error, :unsupported_json_escape}

  defp scan_json_payload(<<?b, _rest::binary>>, true, true),
    do: {:error, :unsupported_json_escape}

  defp scan_json_payload(<<?f, _rest::binary>>, true, true),
    do: {:error, :unsupported_json_escape}

  defp scan_json_payload(<<?n, _rest::binary>>, true, true),
    do: {:error, :unsupported_json_escape}

  defp scan_json_payload(<<?r, _rest::binary>>, true, true),
    do: {:error, :unsupported_json_escape}

  defp scan_json_payload(<<?t, _rest::binary>>, true, true),
    do: {:error, :unsupported_json_escape}

  defp scan_json_payload(<<_byte::8, _rest::binary>>, true, true),
    do: {:error, :unsupported_json_escape}

  defp valid_hex_byte?(byte) when byte >= ?0 and byte <= ?9, do: true
  defp valid_hex_byte?(byte) when byte >= ?a and byte <= ?f, do: true
  defp valid_hex_byte?(byte) when byte >= ?A and byte <= ?F, do: true
  defp valid_hex_byte?(_byte), do: false

  defp json_unicode_escape_codepoint(a, b, c, d) do
    hex_digit_to_int(a) * 16 * 16 * 16 +
      hex_digit_to_int(b) * 16 * 16 +
      hex_digit_to_int(c) * 16 +
      hex_digit_to_int(d)
  end

  defp hex_digit_to_int(byte) when byte >= ?0 and byte <= ?9, do: byte - ?0
  defp hex_digit_to_int(byte) when byte >= ?a and byte <= ?f, do: byte - ?a + 10
  defp hex_digit_to_int(byte) when byte >= ?A and byte <= ?F, do: byte - ?A + 10
  defp hex_digit_to_int(_byte), do: 0

  @doc """
  Parse binary message to Elixir tuple, if message contains event it will be returned as specific
  `Nostr.Event.t()` struct dependent of type of Event
  """
  @spec parse_specific(String.t()) :: t() | struct()
  def parse_specific(msg) when is_binary(msg) do
    msg
    |> JSON.decode!()
    |> do_parse(:specific)
  end

  # Client to relay
  defp do_parse(["EVENT", event], :general) when is_map(event) do
    case Nostr.Event.parse(event) do
      nil -> :error
      parsed -> {:event, parsed}
    end
  end

  defp do_parse(["EVENT", event], :specific) when is_map(event) do
    case Nostr.Event.parse_specific(event) do
      nil -> :error
      parsed -> {:event, parsed}
    end
  end

  defp do_parse(["REQ", sub_id | filters], _type)
       when is_binary(sub_id) and filters != [] do
    parsed = Enum.map(filters, &Nostr.Filter.parse/1)
    {:req, sub_id, parsed}
  end

  defp do_parse(["CLOSE", sub_id], _type) when is_binary(sub_id) do
    {:close, sub_id}
  end

  defp do_parse(["NEG-OPEN", sub_id, filter, initial_message], _type)
       when is_binary(sub_id) and is_map(filter) and is_binary(initial_message) do
    if valid_hex_string?(initial_message) do
      {:neg_open, sub_id, Nostr.Filter.parse(filter), initial_message}
    else
      :error
    end
  end

  defp do_parse(["NEG-MSG", sub_id, message], _type)
       when is_binary(sub_id) and is_binary(message) do
    if valid_hex_string?(message) do
      {:neg_msg, sub_id, message}
    else
      :error
    end
  end

  defp do_parse(["NEG-CLOSE", sub_id], _type) when is_binary(sub_id) do
    {:neg_close, sub_id}
  end

  defp do_parse(["AUTH", event], :general) when is_map(event) do
    case Nostr.Event.parse(event) do
      nil -> :error
      parsed -> {:auth, parsed}
    end
  end

  defp do_parse(["AUTH", event], :specific) when is_map(event) do
    case Nostr.Event.parse_specific(event) do
      nil -> :error
      parsed -> {:auth, parsed}
    end
  end

  # Relay to client
  defp do_parse(["EVENT", sub_id, event], :general) when is_binary(sub_id) and is_map(event) do
    case Nostr.Event.parse(event) do
      nil -> :error
      parsed -> {:event, sub_id, parsed}
    end
  end

  defp do_parse(["EVENT", sub_id, event], :specific) when is_binary(sub_id) and is_map(event) do
    case Nostr.Event.parse_specific(event) do
      nil -> :error
      parsed -> {:event, sub_id, parsed}
    end
  end

  defp do_parse(["NOTICE", message], _type) when is_binary(message) do
    {:notice, message}
  end

  defp do_parse(["EOSE", sub_id], _type) when is_binary(sub_id) do
    {:eose, sub_id}
  end

  defp do_parse(["OK", event_id, success?, message], _type)
       when is_binary(event_id) and is_boolean(success?) and is_binary(message) do
    {:ok, event_id, success?, message}
  end

  defp do_parse(["AUTH", sub_id], _type) when is_binary(sub_id) do
    {:auth, sub_id}
  end

  defp do_parse(["CLOSED", sub_id, message], _type) when is_binary(sub_id) do
    {:closed, sub_id, message}
  end

  defp do_parse(["NEG-ERR", sub_id, reason], _type)
       when is_binary(sub_id) and is_binary(reason) do
    {:neg_err, sub_id, reason}
  end

  defp do_parse(["COUNT", sub_id, payload], _type)
       when is_binary(sub_id) and is_map(payload) and :erlang.is_map_key("count", payload) do
    case parse_count_payload(payload) do
      {:ok, count_payload} -> {:count, sub_id, count_payload}
      :error -> :error
    end
  end

  defp do_parse(["COUNT", sub_id | filters], _type)
       when is_binary(sub_id) and filters != [] do
    case parse_count_filters(filters) do
      {:ok, parsed_filters} -> {:count, sub_id, parsed_filters}
      :error -> :error
    end
  end

  defp do_parse(message, _type) do
    Logger.warning("Parsing unknown message: #{inspect(message)}")
    :error
  end

  defp parse_count_payload(%{"count" => count} = payload) when is_integer(count) do
    with :ok <- validate_approximate(payload),
         :ok <- validate_hll(payload) do
      result = %{count: count}
      result = maybe_put_optional(result, :approximate, payload, "approximate")
      result = maybe_put_optional(result, :hll, payload, "hll")
      {:ok, result}
    else
      :error -> :error
    end
  end

  defp parse_count_payload(_payload), do: :error

  defp parse_count_filters(filters) do
    filters
    |> Enum.reduce_while({:ok, []}, fn
      %{"count" => _count}, _acc -> {:halt, :error}
      filter, {:ok, acc} when is_map(filter) -> {:cont, {:ok, [Nostr.Filter.parse(filter) | acc]}}
      _filter, _acc -> {:halt, :error}
    end)
    |> case do
      {:ok, parsed_filters} -> {:ok, Enum.reverse(parsed_filters)}
      :error -> :error
    end
  end

  defp validate_approximate(payload) do
    case Map.fetch(payload, "approximate") do
      :error -> :ok
      {:ok, value} when is_boolean(value) -> :ok
      {:ok, _value} -> :error
    end
  end

  defp validate_hll(payload) do
    case Map.fetch(payload, "hll") do
      :error -> :ok
      {:ok, value} when is_binary(value) -> if(valid_hll_hex?(value), do: :ok, else: :error)
      {:ok, _value} -> :error
    end
  end

  defp valid_hll_hex?(value) when byte_size(value) == 512 do
    case Base.decode16(value, case: :mixed) do
      {:ok, decoded} -> byte_size(decoded) == 256
      :error -> false
    end
  end

  defp valid_hll_hex?(_value), do: false

  defp valid_hex_string?(value) when is_binary(value) and rem(byte_size(value), 2) == 0 do
    case Base.decode16(value, case: :mixed) do
      {:ok, _decoded} -> true
      :error -> false
    end
  end

  defp valid_hex_string?(_value), do: false

  defp maybe_put_optional(acc, key, payload, payload_key) do
    case Map.fetch(payload, payload_key) do
      {:ok, value} -> Map.put(acc, key, value)
      :error -> acc
    end
  end
end