require 'power_assert/configuration'
require 'power_assert/enable_tracepoint_events'
require 'power_assert/inspector'
require 'power_assert/parser'

module PowerAssert
  class Context
    Value = Struct.new(:name, :value, :column)

    attr_reader :message_proc

    def initialize(base_caller_length)
      @fired = false
      @target_thread = Thread.current
      method_id_set = nil
      return_values = []
      trace_alias_method = PowerAssert.configuration._trace_alias_method
      @trace_return = TracePoint.new(:return, :c_return) do |tp|
        method_id_set ||= @parser.method_id_set
        method_id = SUPPORT_ALIAS_METHOD                      ? tp.callee_id :
                    trace_alias_method && tp.event == :return ? tp.binding.eval('::Kernel.__callee__') :
                                                                tp.method_id
        next if ! method_id_set[method_id]
        next if tp.event == :c_return and
                not (@parser.lineno == tp.lineno and @parser.path == tp.path)
        next unless tp.binding # workaround for ruby 2.2
        locs = PowerAssert.app_caller_locations
        diff = locs.length - base_caller_length
        if (tp.event == :c_return && diff == 1 || tp.event == :return && diff <= 2) and Thread.current == @target_thread
          idx = -(base_caller_length + 1)
          if @parser.path == locs[idx].path and @parser.lineno == locs[idx].lineno
            val = PowerAssert.configuration.lazy_inspection ?
              tp.return_value :
              InspectedValue.new(SafeInspectable.new(tp.return_value).inspect)
            return_values << Value[method_id.to_s, val, nil]
          end
        end
      end
      @message_proc = -> {
        raise RuntimeError, 'call #yield or #enable at first' unless fired?
        @message ||= build_assertion_message(@parser.line, @parser.idents, @parser.binding, return_values).freeze
      }
    end

    def message
      @message_proc.()
    end

    private

    def fired?
      @fired
    end

    def build_assertion_message(line, idents, proc_binding, return_values)
      if PowerAssert.configuration._colorize_message
        line = Pry::Code.new(line).highlighted
      end

      path = detect_path(idents, return_values)
      return line unless path

      delete_unidentified_calls(return_values, path)
      methods, refs = path.partition {|i| i.type == :method }
      return_values.zip(methods) do |i, j|
        unless i.name == j.name
          warn "power_assert: [BUG] Failed to get column: #{i.name}"
          return line
        end
        i.column = j.column
      end
      ref_values = refs.map {|i| Value[i.name, proc_binding.eval(i.name), i.column] }
      vals = (return_values + ref_values).find_all(&:column).sort_by(&:column).reverse
      return line if vals.empty?

      fmt = (0..vals[0].column).map {|i| vals.find {|v| v.column == i } ? "%<#{i}>s" : ' '  }.join
      lines = []
      lines << line.chomp
      lines << sprintf(fmt, vals.each_with_object({}) {|v, h| h[v.column.to_s.to_sym] = '|' }).chomp
      vals.each do |i|
        inspected_val = SafeInspectable.new(Formatter.new(i.value, i.column)).inspect
        inspected_val.each_line do |l|
          map_to = vals.each_with_object({}) do |j, h|
            h[j.column.to_s.to_sym] = [l, '|', ' '][i.column <=> j.column]
          end
          lines << encoding_safe_rstrip(sprintf(fmt, map_to))
        end
      end
      lines.join("\n")
    end

    def detect_path(idents, return_values)
      return @parser.call_paths.flatten.uniq if @parser.method_id_set.empty?
      all_paths = @parser.call_paths
      return_value_names = return_values.map(&:name)
      uniq_calls = uniq_calls(all_paths)
      uniq_call = return_value_names.find {|i| uniq_calls.include?(i) }
      detected_paths = all_paths.find_all do |path|
        method_names = path.find_all {|ident| ident.type == :method }.map(&:name)
        break [path] if uniq_call and method_names.include?(uniq_call)
        return_value_names == method_names
      end
      return nil unless detected_paths.length == 1
      detected_paths[0]
    end

    def uniq_calls(paths)
      all_calls = enum_count_by(paths.map {|path| path.find_all {|ident| ident.type == :method }.map(&:name).uniq }.flatten) {|i| i }
      all_calls.find_all {|_, call_count| call_count == 1 }.map {|name, _| name }
    end

    def delete_unidentified_calls(return_values, path)
      return_value_num_of_calls = enum_count_by(return_values, &:name)
      path_num_of_calls = enum_count_by(path.find_all {|ident| ident.type == :method }, &:name)
      identified_calls = return_value_num_of_calls.find_all {|name, num| path_num_of_calls[name] == num }.map(&:first)
      return_values.delete_if {|val| ! identified_calls.include?(val.name) }
      path.delete_if {|ident| ident.type == :method and ! identified_calls.include?(ident.name) }
    end

    def enum_count_by(enum, &blk)
      Hash[enum.group_by(&blk).map{|k, v| [k, v.length] }]
    end

    def encoding_safe_rstrip(str)
      str.rstrip
    rescue ArgumentError, Encoding::CompatibilityError
      enc = str.encoding
      if enc.ascii_compatible?
        str.b.rstrip.force_encoding(enc)
      else
        str
      end
    end
  end
  private_constant :Context

  class BlockContext < Context
    def initialize(assertion_proc_or_source, assertion_method, source_binding)
      super(0)
      if assertion_proc_or_source.respond_to?(:to_proc)
        @assertion_proc = assertion_proc_or_source.to_proc
        line = nil
      else
        @assertion_proc = source_binding.eval "Proc.new {#{assertion_proc_or_source}}"
        line = assertion_proc_or_source
      end
      @parser = Parser::DUMMY
      @trace_call = TracePoint.new(:call, :c_call) do |tp|
        if PowerAssert.app_context? and Thread.current == @target_thread
          @trace_call.disable
          locs = PowerAssert.app_caller_locations
          path = locs.last.path
          lineno = locs.last.lineno
          line ||= open(path).each_line.drop(lineno - 1).first
          @parser = Parser.new(line, path, lineno, @assertion_proc.binding, assertion_method.to_s)
        end
      end
    end

    def yield
      @fired = true
      invoke_yield(&@assertion_proc)
    end

    private

    def invoke_yield
      @trace_return.enable do
        @trace_call.enable do
          yield
        end
      end
    end
  end
  private_constant :BlockContext

  class TraceContext < Context
    def initialize(binding)
      target_frame, *base = PowerAssert.app_caller_locations
      super(base.length)
      path = target_frame.path
      lineno = target_frame.lineno
      line = open(path).each_line.drop(lineno - 1).first
      @parser = Parser.new(line, path, lineno, binding)
    end

    def enable
      @fired = true
      @trace_return.enable
    end

    def disable
      @trace_return.disable
    end

    def enabled?
      @trace_return.enabled?
    end
  end
  private_constant :TraceContext
end