defmodule Snarkey.Plug do
  @moduledoc """
  A Plug for integrating Snarkey ZKP authentication into Plug-based apps.

  ## Usage

      plug Snarkey.Plug,
        param_key: "proof",
        lookup: &MyApp.Accounts.get_by_blind_id/1,
        fallback_handler: {MyApp.FallbackController, :challenge}

  Expects proof params under `conn.body_params[param_key]` containing `"r"`, `"s"`,
  and optionally `"timestamp"`. The `:lookup` function receives `blind_identifier`
  and must return `%{public_y: binary()}` or `nil`.

  Proof values (`"r"`, `"s"`, `"timestamp"`) are expected as decimal strings and
  are converted to raw binary-encoded integers before calling `Snarkey.authenticate/3`.

  On success, assigns `:snarkey_auth => {:ok, :authenticated}` to `conn.assigns`.
  On fallback, invokes the fallback handler with `conn` and the challenge, or
  assigns `:snarkey_auth => {:fallback, challenge}` if no handler is configured.
  On error, returns a 401 response.
  """

  @defaults %{param_key: "proof", lookup: nil, fallback_handler: nil}

  def init(opts) do
    Map.merge(@defaults, Map.new(opts))
  end

  def call(conn, opts) do
    blind_identifier = conn.private.blind_identifier

    with proof when is_map(proof) <- Map.get(conn.params, opts.param_key),
         user when not is_nil(user) <- opts.lookup.(blind_identifier) do
      do_auth(conn, opts, user, proof)
    else
      _ -> unauthorized(conn)
    end
  end

  defp do_auth(conn, opts, user, proof) do
    proof =
      proof
      |> Map.put_new("timestamp", Integer.to_string(System.system_time(:second)))
      |> encode_proof_values()

    case Snarkey.authenticate(user.public_y, atomize_keys(proof),
           user_id: conn.private.blind_identifier
         ) do
      {:ok, :authenticated} ->
        Plug.Conn.assign(conn, :snarkey_auth, {:ok, :authenticated})

      {:fallback, challenge} ->
        handle_fallback(conn, opts, challenge)

      {:error, _reason} ->
        unauthorized(conn)
    end
  end

  defp encode_proof_values(proof) do
    Map.new(proof, fn
      {"r", val} -> {"r", decimal_to_binary(val)}
      {"s", val} -> {"s", decimal_to_binary(val)}
      {"timestamp", val} -> {"timestamp", decimal_to_binary(val)}
      {k, v} -> {k, v}
    end)
  end

  defp decimal_to_binary(val) when is_integer(val), do: :binary.encode_unsigned(val)

  defp decimal_to_binary(val) when is_binary(val),
    do: val |> String.to_integer() |> :binary.encode_unsigned()

  defp handle_fallback(conn, %{fallback_handler: {mod, fun}}, challenge) do
    apply(mod, fun, [conn, challenge])
  end

  defp handle_fallback(conn, _opts, challenge) do
    Plug.Conn.assign(conn, :snarkey_auth, {:fallback, challenge})
  end

  defp unauthorized(conn) do
    conn
    |> Plug.Conn.put_resp_content_type("application/json")
    |> Plug.Conn.send_resp(401, ~s({"error":"unauthorized"}))
  end

  defp atomize_keys(map) do
    Map.new(map, fn {key, val} -> {String.to_existing_atom(key), val} end)
  rescue
    ArgumentError ->
      Map.new(map, fn {key, val} -> {String.to_atom(key), val} end)
  end
end