defmodule BallsPDS.JWT do
  alias BallsPDS.WAC
  require Logger
  use Joken.Config

  def generate_my_jwk() do
    raw_private_key =
      Base.decode16!(Application.get_env(:balls_pds, :owner_private_key), case: :lower)

    generate_jwk(raw_private_key)
  end

  def generate_jwk(<<raw_private_key::binary-size(32)>>) do
    public_key = :crypto.generate_key(:eddsa, :ed25519, raw_private_key) |> elem(0)

    %{
      "kty" => "OKP",
      "crv" => "Ed25519",
      "d" => Base.url_encode64(raw_private_key, padding: false),
      "x" => Base.url_encode64(public_key, padding: false)
    }
  end

  def query_public_jwk(ap_id, key_id) do
    case WAC.cached_query_public_key(ap_id, key_id) do
      {:ok, raw_public_key} ->
        encoded_key = Base.url_encode64(raw_public_key, padding: false)

        {:ok,
         %{
           # Key Type: Octet Key Pair
           "kty" => "OKP",
           # Curve: Ed25519
           "crv" => "Ed25519",
           # Public Key (Base64Url-encoded)
           "x" => encoded_key
         }}
    end
  end

  defp get_kid(jwt) when is_binary(jwt) do
    with {:kid, %{fields: %{"kid" => kid}}} <- {:kid, JOSE.JWT.peek_protected(jwt)} do
      kid
    else
      {:kid, protected} ->
        Logger.debug("protected: #{inspect(protected)}")
        nil

      error ->
        Logger.error(inspect(error))
        nil
    end
  end

  def extract_key_info(jwt) when is_binary(jwt) do
    with {:subject, {:ok, %{"sub" => subject}}} <- {:subject, JOSE.JWT.peek_payload(jwt)},
         {:kid, kid} <- {:kid, get_kid(jwt)} do
      {:ok, %{subject: subject, id: kid}}
    else
      {err, {:error, error}} ->
        Logger.error("extracting key info from JWT: #{err}: #{inspect(error)}")
        {:error, error}
    end
  end

  def generate_jwt(jwk, days) when is_integer(days) and days > 0 do
    Logger.debug("Creating signer.")
    signer = Joken.Signer.create("EdDSA", jwk, %{"kid" => "key-1"})

    id = Application.get_env(:balls_pds, :owner_ap_id)

    claims =
      default_claims(%{
        "aud" => Application.get_env(:balls_pds, :owner_ap_id),
        "iat" => DateTime.utc_now() |> DateTime.to_unix(),
        "exp" => DateTime.utc_now() |> DateTime.add(30, :day) |> DateTime.to_unix()
      })

    additional_claims = %{"sub" => id}

    Logger.debug("Signing with claims: #{inspect(claims)}")
    Joken.generate_and_sign!(claims, additional_claims, signer)
  end

  def generate_jwt(days) when is_integer(days) and days > 0 do
    jwk = generate_my_jwk()
    generate_jwt(jwk, days)
  end

  def verify_jwt(jwt, jwk) do
    public_jwk = Map.drop(jwk, ["d"])
    signer = Joken.Signer.create("EdDSA", public_jwk)

    case Joken.verify_and_validate(public_jwk, jwt, signer) do
      {:ok, claims} -> {:ok, claims}
      {:error, reason} -> {:error, reason}
    end
  end

  def verify_jwt(jwt, jwk, subject) when is_binary(subject) do
    case verify_jwt(jwt, jwk) do
      {:ok, claims = %{"sub" => ^subject}} -> {:ok, claims}
      {:ok, %{"sub" => _wrong_subject}} -> {:error, :wrong_subject}
      {:ok, _claims} -> {:error, :missing_subject}
      error = {:error, _} -> error
    end
  end
end