document_utils.rb

# frozen_string_literal: true

require 'base64'
require 'time'

module Smithy
  module Schema
    # @api private
    # Document Utilities to help (de)construct data to/from Smithy document
    module DocumentUtils
      class << self
        # Used to transform untyped data
        def format(data)
          return if data.nil?

          case data
          when Time
            data.to_i # timestamp format is "epoch-seconds" by default
          when Hash
            data.transform_values { |v| format(v) }
          when Array
            data.map { |d| format(d) }
          else
            data
          end
        end

        # Used to apply data to runtime shape
        def apply(data, schema, type = nil)
          case shape(schema)
          when Shapes::StructureShape then apply_structure(data, schema, type)
          when Shapes::UnionShape then apply_union(data, schema, type)
          when Shapes::ListShape then apply_list(data, schema)
          when Shapes::MapShape then apply_map(data, schema)
          when Shapes::TimestampShape then apply_timestamp(data, schema)
          when Shapes::BlobShape then Base64.decode64(data)
          else data
          end
        end

        # rubocop:disable Metrics/CyclomaticComplexity
        def extract(data, schema, opts = {})
          return if data.nil?

          case shape(schema)
          when Shapes::StructureShape then extract_structure(data, schema, opts)
          when Shapes::UnionShape     then extract_union(data, schema, opts)
          when Shapes::ListShape      then extract_list(data, schema)
          when Shapes::MapShape       then extract_map(data, schema)
          when Shapes::BlobShape      then extract_blob(data)
          when Shapes::TimestampShape then extract_timestamp(data, schema, opts)
          else data
          end
        end
        # rubocop:enable Metrics/CyclomaticComplexity

        private

        def apply_list(data, schema)
          shape = shape(schema)
          data.map do |v|
            next if v.nil?

            apply(v, shape.member)
          end
        end

        def apply_map(data, schema)
          shape = shape(schema)
          data.transform_values do |v|
            if v.nil?
              nil
            else
              apply(v, shape.value)
            end
          end
        end

        def apply_structure(data, schema, type)
          shape = shape(schema)

          type = shape.type.new if type.nil?
          data.each do |k, v|
            name =
              if (member = member_with_json_name(k, shape))
                shape.name_by_member_name(member.name)
              else
                member_name(shape, k)
              end
            next if name.nil?

            type[name] = apply(v, shape.member(name))
          end
          type
        end

        def apply_timestamp(data, schema)
          data = data.is_a?(Numeric) ? Time.at(data) : Time.parse(data)
          time(data, timestamp_format(schema))
        end

        def apply_union(data, schema, type)
          shape = shape(schema)
          key, value = data.flatten
          return if key.nil?

          if (member = member_with_json_name(key, shape))
            apply_union_member(member.name, value, shape, type)
          elsif shape.name_by_member_name?(key)
            apply_union_member(key, value, shape, type)
          else
            shape.member_type(:unknown).new(key, value)
          end
        end

        def apply_union_member(key, value, shape, type)
          member_name = shape.name_by_member_name(key)
          type = shape.member_type(member_name) if type.nil?
          type.new(apply(value, shape.member(member_name)))
        end

        def extract_blob(data)
          Base64.strict_encode64(data.is_a?(String) ? data : data.read)
        end

        def extract_list(data, schema)
          shape = shape(schema)
          data.collect { |v| extract(v, shape.member) }
        end

        def extract_map(data, schema)
          shape = shape(schema)
          data.each.with_object({}) { |(k, v), h| h[k] = extract(v, shape.value) }
        end

        def extract_structure(data, schema, opts)
          shape = shape(schema)
          data.to_h.each_with_object({}) do |(k, v), o|
            next unless shape.member?(k)

            member_shape = shape.member(k)
            member_name = resolve_member_name(member_shape, opts)
            o[member_name] = extract(v, member_shape, opts)
          end
        end

        def extract_timestamp(data, schema, opts)
          return unless data.is_a?(Time)

          trait = timestamp_format(schema) if opts[:use_timestamp_format]
          time(data, trait)
        end

        # rubocop:disable Metrics/AbcSize
        def extract_union(data, schema, opts)
          h = {}
          shape = shape(schema)
          if data.is_a?(Schema::Union)
            member_shape = shape.member_by_type(data.class)
            member_name = resolve_member_name(member_shape, opts)
            h[member_name] = extract(data, member_shape).value
          else
            key, value = data.first
            if shape.member?(key)
              member_shape = shape.member(key)
              member_name = resolve_member_name(member_shape, opts)
              h[member_name] = extract(value, member_shape)
            end
          end
          h
        end
        # rubocop:enable Metrics/AbcSize

        def member_name(schema, key)
          return unless schema.name_by_member_name?(key) || schema.member?(key.to_sym)

          schema.name_by_member_name(key) || key.to_sym
        end

        def member_with_json_name(name, shape)
          shape.members.values.find do |v|
            v.traits['smithy.api#jsonName'] == name if v.traits.include?('smithy.api#jsonName')
          end
        end

        def resolve_member_name(member_shape, opts)
          if opts[:use_json_name] && member_shape.traits['smithy.api#jsonName']
            member_shape.traits['smithy.api#jsonName']
          else
            member_shape.name
          end
        end

        def shape(schema)
          schema.is_a?(Shapes::MemberShape) ? schema.shape : schema
        end

        # The following steps are taken to determine the format of timestamp:
        # Use the timestampFormat trait of the member, if present.
        # Use the timestampFormat trait of the shape, if present.
        # If none of the above applies, use epoch-seconds as default
        def timestamp_format(schema)
          if schema.traits['smithy.api#timestampFormat']
            schema.traits['smithy.api#timestampFormat']
          elsif schema.shape.traits['smithy.api#timestampFormat']
            schema.shape.traits['smithy.api#timestampFormat']
          else
            'epoch-seconds'
          end
        end

        def time(data, trait = nil)
          if trait
            case trait
            when 'http-date'
              data.utc.iso8601
            when 'date-time'
              data.utc.httpdate
            when 'epoch-seconds'
              data.utc.to_i
            else
              raise "unhandled timestamp format `#{value}`"
            end
          else
            data.utc.to_i # default format
          end
        end
      end
    end
  end
end