# frozen_string_literal: true

module IpynbDiff

  # Transforms Jupyter output data into markdown
  class OutputTransformer

    ORDERED_KEYS = {
      'execute_result' => %w[image/png image/svg+xml image/jpeg text/markdown text/latex text/plain],
      'display_data' => %w[image/png image/svg+xml image/jpeg text/markdown text/latex]
    }.freeze

    def transform(output)
      case (output_type = output['output_type'])
      when 'error'
        transform_error(output['traceback'])
      when 'execute_result', 'display_data'
        transform_non_error(ORDERED_KEYS[output_type], output['data'])
      end
    end

    def transform_error(traceback)
      traceback.map do |t|
        t.split("\n").map do |line|
          line.gsub(/\[[0-9][0-9;]*m/, '').sub("\u001B", '    ').gsub(/\u001B/, '').rstrip << "\n"
        end
      end
    end

    def transform_non_error(accepted_keys, elements)
      accepted_keys.map do |key|
        transform_element(key, elements[key]) if elements.key?(key)
      end.flatten
    end

    def transform_element(output_type, output_element)
      case output_type
      when 'image/png', 'image/jpeg'
        transform_image(output_type, output_element)
      when 'image/svg+xml'
        transform_svg(output_element)
      when 'text/markdown', 'text/latex', 'text/plain'
        transform_text(output_element)
      end
    end

    def transform_image(image_type, image_content)
      ["    ![](data:#{image_type};base64,#{image_content.gsub("\n", '')})", "\n"]
    end

    def transform_svg(image_content)
      lines = image_content.is_a?(Array) ? image_content : [image_content]

      single_line = lines.map(&:strip).join('').gsub(/\s+/, ' ')

      ["    ![](data:image/svg+xml;utf8,#{single_line})", "\n"]
    end

    def transform_text(text_content)
      lines = text_content.is_a?(Array) ? text_content : [text_content]

      lines.map { |line| "    #{line}" }.append("\n")
    end
  end
end
