class JSON::JWT

def as_json(options = {})
  case options[:syntax]
  when :general
    {
      payload: Base64.urlsafe_encode64(self.to_json, padding: false),
      signatures: [{
        protected: Base64.urlsafe_encode64(header.to_json, padding: false),
        signature: Base64.urlsafe_encode64(signature.to_s, padding: false)
      }]
    }
  when :flattened
    {
      protected: Base64.urlsafe_encode64(header.to_json, padding: false),
      payload:   Base64.urlsafe_encode64(self.to_json, padding: false),
      signature: Base64.urlsafe_encode64(signature.to_s, padding: false)
    }
  else
    super
  end
end

def decode_compact_serialized(jwt_string, key_or_secret, algorithms = nil, encryption_methods = nil, allow_blank_payload = false)
  case jwt_string.count('.') + 1
  when JWS::NUM_OF_SEGMENTS
    JWS.decode_compact_serialized jwt_string, key_or_secret, algorithms, allow_blank_payload
  when JWE::NUM_OF_SEGMENTS
    if allow_blank_payload
      raise InvalidFormat.new("JWE w/ blank payload is not supported.")
    else
      JWE.decode_compact_serialized jwt_string, key_or_secret, algorithms, encryption_methods
    end
  else
    raise InvalidFormat.new("Invalid JWT Format. JWT should include #{JWS::NUM_OF_SEGMENTS} or #{JWE::NUM_OF_SEGMENTS} segments.")
  end
end

def decode_json_serialized(input, key_or_secret, algorithms = nil, encryption_methods = nil, allow_blank_payload = false)
  input = input.with_indifferent_access
  if (input[:signatures] || input[:signature]).present?
    JWS.decode_json_serialized input, key_or_secret, algorithms, allow_blank_payload
  elsif input[:ciphertext].present?
    JWE.decode_json_serialized input, key_or_secret, algorithms, encryption_methods
  else
    raise InvalidFormat.new("Unexpected JOSE JSON Serialization Format.")
  end
end

def encrypt(public_key_or_secret, algorithm = :RSA1_5, encryption_method = :'A128CBC-HS256')
  jwe = JWE.new self
  jwe.kid ||= public_key_or_secret[:kid] if public_key_or_secret.is_a? JSON::JWK
  jwe.alg = algorithm
  jwe.enc = encryption_method
  jwe.encrypt! public_key_or_secret
end

def initialize(claims = {})
  @content_type = 'application/jwt'
  self.typ = :JWT
  self.alg = :none
  update claims
  unless claims.nil?
    [:exp, :nbf, :iat].each do |key|
      self[key] = self[key].to_i if self[key]
    end
  end
end

def pretty_generate
  [
    JSON.pretty_generate(header),
    JSON.pretty_generate(self)
  ]
end

def pretty_generate(jwt_string)
  decode(jwt_string, :skip_verification).pretty_generate
end

def sign(private_key_or_secret, algorithm = :autodetect)
  jws = JWS.new self
  jws.kid ||= private_key_or_secret[:kid] if private_key_or_secret.is_a? JSON::JWK
  jws.alg = algorithm
  jws.sign! private_key_or_secret
end

def to_json *args
  if @blank_payload && args.empty?
    ''
  else
    super
  end
end

def to_s
  [
    header.to_json,
    self.to_json,
    signature
  ].collect do |segment|
    Base64.urlsafe_encode64 segment.to_s, padding: false
  end.join('.')
end

def update claims
  if claims.nil?
    @blank_payload = true
  else
    super
  end
end