defmodule BallsPDS.JWT do
  alias BallsPDS.WAC
  require Logger

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

  def generate_jwk(raw_private_key) when is_binary(raw_private_key) do
    private_key = raw_private_key |> Base.encode64()
    {:ok, public_key} = :crypto.generate_key(:eddsa, :ed25519, private_key)
    JOSE.JWK.from_key({:okp, :Ed25519, public_key, private_key})
  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

  def extract_key_info(jwt) when is_binary(jwt) do
    with {:subject, {:ok, %{"sub" => subject}}} <- {:subject, JOSE.JWT.peek_payload(jwt)},
         {:kid, {:ok, %{"kid" => kid}}} <- {:kid, JOSE.JWT.peek_protected(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(days \\ 30) when is_integer(days) and days > 0 do
    jwk = generate_my_jwk()
    signer = Joken.Signer.create("EdDSA", jwk)

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

    claims = %{
      "iss" => id,
      "sub" => id,
      "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()
    }

    Joken.generate_and_sign!(claims, signer)
  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(signer, jwt) do
      {:ok, claims} -> {:ok, claims}
      {:error, reason} -> {:error, reason}
    end
  end
end