require 'terminal-table'
require 'fastlane_core/print_table'

module Fastlane
  class PrettyTable
    def self.print(data:, columns:, title: nil, column_names: nil)
      if data.is_a?(Hash)
        print_rows(
          data: data.filter { |k, v| columns.include?(k) }.map { |k, v| { key: (column_names ? column_names[k] : nil) || k, value: v } },
          columns: [:key, :value],
          title: title,
          header: false
        )
      else
        print_rows(
          data: data,
          columns: columns,
          title: title,
          column_names: column_names,
          header: true
        )
      end
    end

    def self.print_rows(data:, columns:, title: nil, column_names: nil, footer_index: nil, header: true)
      headings = header ? columns.map do |column_key|
        if column_names.nil? || column_names[column_key].nil?
          column_key.to_s.split("_").map(&:capitalize).join(" ")
        else
          column_names[column_key]
        end
      end : nil

      rows = data.map do |row|
        column_values = columns.map do |column_key|
          if row[column_key].is_a?(Integer)
            row[column_key].to_s.yellow
          elsif row[column_key].is_a?(Array)
            row[column_key].count.to_s.green
          elsif column_key == columns.first
            row[column_key].to_s.cyan
          else
            row[column_key] || ""
          end
        end
        column_values
      end

      transformed_rows = FastlaneCore::PrintTable.transform_output(rows)
      table = Terminal::Table.new(
        headings: headings,
        title: title.green,
        rows: footer_index.nil? ? transformed_rows : [*transformed_rows.slice(0, footer_index), :separator, *transformed_rows.slice(footer_index..-1)]
      )

      puts("")
      puts(table)
      puts("")
    end
  end
end