class Pry::Code

def ==(other)
  if other.is_a?(Code)
    other_lines = other.instance_variable_get(:@lines)
    @lines.each_with_index.all? { |loc, i| loc == other_lines[i] }
  else
    to_s.chomp == other.to_s.chomp
  end
end

def after(lineno, lines = 1)
  return self unless lineno
  select do |loc|
    loc.lineno > lineno && loc.lineno <= lineno + lines
  end
end

def alter(&block)
  dup.tap { |o| o.instance_eval(&block) }
end

def around(lineno, lines = 1)
  return self unless lineno
  select do |loc|
    loc.lineno >= lineno - lines && loc.lineno <= lineno + lines
  end
end

def before(lineno, lines = 1)
  return self unless lineno
  select do |loc|
    loc.lineno >= lineno - lines && loc.lineno < lineno
  end
end

def between(start_line, end_line = nil)
  return self unless start_line
  code_range = CodeRange.new(start_line, end_line)
  alter do
    @lines = @lines[code_range.indices_range(@lines)] || []
  end
end

def comment_describing(line_number)
  self.class.comment_describing(raw, line_number)
end

def expression_at(line_number, consume = 0)
  self.class.expression_at(raw, line_number, consume: consume)
end

def from_file(filename, code_type = nil)
  code_file = CodeFile.new(filename, code_type)
  new(code_file.code, 1, code_file.code_type)
end

def from_method(meth, start_line = nil)
  meth = Pry::Method(meth)
  start_line ||= meth.source_line || 1
  new(meth.source, start_line, meth.source_type)
end

def from_module(mod, candidate_rank = 0, start_line = nil)
  candidate = Pry::WrappedModule(mod).candidate(candidate_rank)
  start_line ||= candidate.line
  new(candidate.source, start_line, :ruby)
end

def grep(pattern)
  return self unless pattern
  pattern = Regexp.new(pattern)
  select do |loc|
    loc.line =~ pattern
  end
end

def highlighted
  print_to_output(''.dup, true)
end

def initialize(lines = [], start_line = 1, code_type = :ruby)
  lines = lines.lines if lines.is_a? String
  @lines = lines.each_with_index.map do |line, lineno|
    LOC.new(line, lineno + start_line.to_i)
  end
  @code_type = code_type
  @with_marker = @with_indentation = @with_line_numbers = nil
end

def length
  @lines ? @lines.length : 0
end

def max_lineno_width
  !@lines.empty? ? @lines.last.lineno.to_s.length : 0
end

def method_missing(method_name, *args, &block)
  if (string = to_s).respond_to?(method_name)
    string.__send__(method_name, *args, &block)
  else
    super
  end
end

def nesting_at(line_number)
  Pry::Indent.nesting_at(raw, line_number)
end

def print_to_output(output, color = false)
  @lines.each do |loc|
    loc = loc.dup
    loc.colorize(@code_type)              if color
    loc.add_line_number(max_lineno_width, color) if @with_line_numbers
    loc.add_marker(@marker_lineno)        if @with_marker
    loc.indent(@indentation_num)          if @with_indentation
    output << loc.line
    output << "\n"
  end
  output
end

def push(line)
  line_number = @lines.any? ? @lines.last.lineno + 1 : 1
  @lines.push(LOC.new(line, line_number))
end

def raw
  @lines.map(&:line).join("\n") << "\n"
end

def reject(&block)
  alter do
    @lines = @lines.reject(&block)
  end
end

def respond_to_missing?(method_name, include_private = false)
  ''.respond_to?(method_name, include_private) || super
end

def select(&block)
  alter do
    @lines = @lines.select(&block)
  end
end

def take_lines(start_line, num_lines)
  start_idx =
    if start_line >= 0
      @lines.index { |loc| loc.lineno >= start_line } || @lines.length
    else
      [@lines.length + start_line, 0].max
    end
  alter do
    @lines = @lines.slice(start_idx, num_lines)
  end
end

def to_s
  print_to_output(''.dup, false)
end

def with_indentation(spaces = 0)
  alter do
    @with_indentation = !!spaces
    @indentation_num  = spaces
  end
end

def with_line_numbers(y_n = true)
  alter do
    @with_line_numbers = y_n
  end
end

def with_marker(lineno = 1)
  alter do
    @with_marker   = !!lineno
    @marker_lineno = lineno
  end
end