class Google::Auth::IDTokens::KeyInfo

def ensure_json_parsed input
  return input unless input.is_a? String
  JSON.parse input
rescue JSON::ParserError
  raise KeySourceError, "Unable to parse JSON"
end

def extract_ec_key jwk
  begin
    x_data = Base64.urlsafe_decode64 jwk[:x]
    y_data = Base64.urlsafe_decode64 jwk[:y]
  rescue ArgumentError
    raise KeySourceError, "Badly formatted key data"
  end
  curve_name = CURVE_NAME_MAP[jwk[:crv]]
  raise KeySourceError, "Unsupported EC curve #{jwk[:crv]}" unless curve_name
  group = OpenSSL::PKey::EC::Group.new curve_name
  x_hex = x_data.unpack1 "H*"
  y_hex = y_data.unpack1 "H*"
  bn = OpenSSL::BN.new ["04#{x_hex}#{y_hex}"].pack("H*"), 2
  point =  OpenSSL::PKey::EC::Point.new group, bn
  sequence = OpenSSL::ASN1::Sequence([
                                       OpenSSL::ASN1::Sequence([OpenSSL::ASN1::ObjectId("id-ecPublicKey"),
                                                                OpenSSL::ASN1::ObjectId(curve_name)]),
                                       OpenSSL::ASN1::BitString(point.to_octet_string(:uncompressed))
                                     ])
  OpenSSL::PKey::EC.new sequence.to_der
end

def extract_rsa_key jwk
  begin
    n_data = Base64.urlsafe_decode64 jwk[:n]
    e_data = Base64.urlsafe_decode64 jwk[:e]
  rescue ArgumentError
    raise KeySourceError, "Badly formatted key data"
  end
  n_bn = OpenSSL::BN.new n_data, 2
  e_bn = OpenSSL::BN.new e_data, 2
  sequence = [OpenSSL::ASN1::Integer.new(n_bn), OpenSSL::ASN1::Integer.new(e_bn)]
  rsa_key =  OpenSSL::PKey::RSA.new OpenSSL::ASN1::Sequence(sequence).to_der
  rsa_key.public_key
end

def from_jwk jwk
  jwk = symbolize_keys ensure_json_parsed jwk
  key = case jwk[:kty]
        when "RSA"
          extract_rsa_key jwk
        when "EC"
          extract_ec_key jwk
        when nil
          raise KeySourceError, "Key type not found"
        else
          raise KeySourceError, "Cannot use key type #{jwk[:kty]}"
        end
  new id: jwk[:kid], key: key, algorithm: jwk[:alg]
end

def from_jwk_set jwk_set
  jwk_set = symbolize_keys ensure_json_parsed jwk_set
  jwks = jwk_set[:keys]
  raise KeySourceError, "No keys found in jwk set" unless jwks
  jwks.map { |jwk| from_jwk jwk }
end

def initialize id: nil, key: nil, algorithm: nil
  @id = id
  @key = key
  @algorithm = algorithm
end

def symbolize_keys hash
  result = {}
  hash.each { |key, val| result[key.to_sym] = val }
  result
end