defmodule Cashubrew.Nuts.Nut00.BDHKE do
  @moduledoc """
  Cryptographic functions for the Cashubrew mint, including BDHKE implementation.
  Check [NUT-00](https://cashubtc.github.io/nuts/00/) for more information.
  Protocol summary:
  - Alice (user) blinds the message and sends it to Bob (mint).
  - Bob signs the blinded message.
  - Alice unblinds the signature.
  - Carol (user) verifies the signature.
  """
  require Logger

  alias Cashubrew.Crypto.Secp256k1Utils
  alias ExSecp256k1

  # secp256k1 parameters
  defp secp256k1_n do
    0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141
  end

  @doc """
  Generate a new keypair.
  Takes private key as input or generates a new one if not provided.
  """
  def generate_keypair(private_key \\ nil) do
    with {:ok, private_key} <- generate_private_key(private_key),
         {:ok, public_key} <- generate_public_key(private_key) do
      {:ok, {private_key, public_key}}
    end
  end

  def generate_private_key(private_key \\ nil) do
    case private_key do
      nil -> {:ok, :crypto.strong_rand_bytes(32)}
      <<0, key::binary-size(32)>> -> {:ok, key}
      <<key::binary-size(32)>> -> {:ok, key}
      _ -> {:error, "Invalid private key format"}
    end
  end

  def generate_public_key(private_key) do
    with {:ok, public_key} <- ExSecp256k1.create_public_key(private_key) do
      ExSecp256k1.public_key_compress(public_key)
    end
  end

  def load_public_key(hash) do
    with {:ok, public_key} <- ExSecp256k1.public_key_decompress(hash) do
      ExSecp256k1.public_key_compress(public_key)
    end
  end

  @doc """
  /// The point is generated by hashing the message with a domain separator and then
  /// iteratively trying to compute a point from the hash. An increasing uint32 counter
  /// (byte order little endian) is appended to the hash until a point is found that lies on the curve.

  /// The chance of finding a valid point is 50% for every iteration. The maximum number of iterations
  /// is 2\\*\\*16. If no valid point is found after 2**16 iterations, a ValueError is raised (this should
  /// never happen in practice).
  """
  def hash_to_curve(message) do
    :crypto.hash(:sha256, "Secp256k1_HashToCurve_Cashu_" <> message)
    |> find_valid_point()
  end

  defp find_valid_point(msg_hash, counter \\ 0)

  defp find_valid_point(msg_hash, counter) when counter < 0x10000 do
    to_hash = msg_hash <> <<counter::little-32>>
    hash = :crypto.hash(:sha256, to_hash)

    case load_public_key(<<2>> <> hash) do
      {:ok, public_key} -> public_key
      _ -> find_valid_point(msg_hash, counter + 1)
    end
  end

  # Not supposed to ever happen
  defp find_valid_point(_, _) do
    raise("No valid point found")
  end

  @doc """
  Alice's step 1: Blind the message.
  Alice is the sending user.
  Alice picks secret x and computes Y = hash_to_curve(x).
  Alice sends to Bob: B_ = Y + rG with r being a random blinding factor (blinding).
  This operation is called blinding.
  """
  def step1_alice(secret_msg, r) do
    y = hash_to_curve(secret_msg)

    with {:ok, r_pub} <- ExSecp256k1.create_public_key(r),
         {:ok, r_pub_compressed} <- ExSecp256k1.public_key_compress(r_pub),
         {:ok, b_prime} <- Secp256k1Utils.point_add(y, r_pub_compressed) do
      # B_ = Y + rG
      {:ok, {b_prime, r}}
    end
  end

  @doc """
  Bob's step 2: Sign the blinded message.
  Bob is the mint.
  This operation is called signing.
  """
  def step2_bob(b_prime, a) do
    with {:ok, c_prime} <- Secp256k1Utils.point_mul(b_prime, a),
         {:ok, {e, s}} <- step2_bob_dleq(b_prime, a) do
      {:ok, {c_prime, e, s}}
    end
  end

  @doc """
  Alice's step 3: Unblind the signature.
  Alice can calculate the unblinded key as C_ - rK = kY + krG - krG = kY = C.
  This operation is called unblinding.
  """
  def step3_alice(c_prime, r, a_pub) do
    with {:ok, r_a_pub} <- Secp256k1Utils.point_mul(a_pub, r) do
      Secp256k1Utils.point_sub(c_prime, r_a_pub)
    end
  end

  @doc """
  Verify the signature.
  Carol is the receiving user.
  This operation is called verification.
  """
  def verify?(a, c, secret_msg) do
    y = hash_to_curve(secret_msg)

    case Secp256k1Utils.point_mul(y, a) do
      {:ok, a_y} -> Secp256k1Utils.point_equal?(c, a_y)
      {:err, _} -> false
    end
  end

  @doc """
  Generate DLEQ proof
  """
  def step2_bob_dleq(b_prime, a, p \\ nil) do
    p = p || :crypto.strong_rand_bytes(32)

    with {:ok, r1} <- ExSecp256k1.create_public_key(p),
         {:ok, r1_compressed} <- ExSecp256k1.public_key_compress(r1),
         {:ok, r2} <- Secp256k1Utils.point_mul(b_prime, p),
         {:ok, a_pub} <- ExSecp256k1.create_public_key(a),
         {:ok, a_pub_compressed} <- ExSecp256k1.public_key_compress(a_pub),
         {:ok, c_prime} <- Secp256k1Utils.point_mul(b_prime, a) do
      e = hash_e(r1_compressed, r2, a_pub_compressed, c_prime)
      e_scalar = :binary.decode_unsigned(e)
      a_times_e = Secp256k1Utils.mod_mul(:binary.decode_unsigned(a), e_scalar, secp256k1_n())
      s = Secp256k1Utils.mod_add(:binary.decode_unsigned(p), a_times_e, secp256k1_n())
      s_bin = :binary.encode_unsigned(s) |> Secp256k1Utils.pad_left(32)
      {:ok, {e, s_bin}}
    end
  end

  @doc """
  Alice verifies DLEQ proof
  """
  def alice_verify_dleq?(b_prime, c_prime, e, s, a_pub) do
    with {:ok, s_g} <- ExSecp256k1.create_public_key(s),
         {:ok, s_g_compressed} <- ExSecp256k1.public_key_compress(s_g),
         {:ok, e_a} <- Secp256k1Utils.point_mul(a_pub, e),
         {:ok, r1} <- Secp256k1Utils.point_sub(s_g_compressed, e_a),
         {:ok, s_b_prime} <- Secp256k1Utils.point_mul(b_prime, s),
         {:ok, e_c_prime} <- Secp256k1Utils.point_mul(c_prime, e),
         {:ok, r2} <- Secp256k1Utils.point_sub(s_b_prime, e_c_prime) do
      hash_e(r1, r2, a_pub, c_prime) == e
    else
      _ -> false
    end
  end

  @doc """
  Carol verifies DLEQ proof
  """
  def carol_verify_dleq?(secret_msg, r, c, e, s, a_pub) do
    y = hash_to_curve(secret_msg)

    with {:ok, r_pub} <- ExSecp256k1.create_public_key(r),
         {:ok, r_pub_compressed} <- ExSecp256k1.public_key_compress(r_pub),
         {:ok, r_a_pub} <- Secp256k1Utils.point_mul(a_pub, r),
         {:ok, c_prime} <- Secp256k1Utils.point_add(c, r_a_pub),
         {:ok, b_prime} <- Secp256k1Utils.point_add(y, r_pub_compressed) do
      alice_verify_dleq?(b_prime, c_prime, e, s, a_pub)
    else
      _ -> false
    end
  end

  @doc """
  Hash function for DLEQ
  """
  def hash_e(r1, r2, a, c_prime) do
    keys = [r1, r2, a, c_prime]

    data =
      Enum.map_join(keys, "", fn key ->
        {:ok, uncompressed} = ExSecp256k1.public_key_decompress(key)
        uncompressed |> Base.encode16(case: :lower)
      end)

    :crypto.hash(:sha256, data)
  end
end
