class ANSIString

# File lib/ansi_string.rb, line 12
def initialize(str)
  process_string raw_string_for(str)
end

# File lib/ansi_string.rb, line 16
def +(other)
  self.class.new @raw + raw_string_for(other)
end

# File lib/ansi_string.rb, line 20
def <<(other)
  range = length..length
  str = replace_in_string(range, other)
  process_string raw_string_for(str)
  self
end

# File lib/ansi_string.rb, line 219
def <=>(other)
  (other.class == self.class && @raw <=> other.raw)
end

# File lib/ansi_string.rb, line 215
def ==(other)
  (other.class == self.class && other.raw == @raw) || (other.kind_of?(String) && other == @raw)
end

# File lib/ansi_string.rb, line 39
def [](range)
  # convert numeric position to a range
  range = (range..range) if range.is_a?(Integer)

  range_begin = range.begin
  range_end = range.end

  if range.exclude_end?
    if range_begin == 0 && range_end == 0
      return ""
    else
      range_end -= 1
    end
  end

  range_begin = @without_ansi.length - range.begin.abs if range.begin < 0
  range_end = @without_ansi.length - range.end.abs if range.end < 0

  str = build_string_with_ansi_for(range_begin..range_end)
  ANSIString.new str if str
end

# File lib/ansi_string.rb, line 61
def []=(range, replacement_str)
  # convert numeric position to a range
  range = (range..range) if range.is_a?(Integer)

  range_begin = range.begin
  range_end = range.exclude_end? ? range.end - 1 : range.end

  range_begin = @without_ansi.length - range.begin.abs if range.begin < 0
  range_end = @without_ansi.length - range.end.abs if range.end < 0

  updated_string = replace_in_string(range_begin..range_end, replacement_str)
  process_string raw_string_for(updated_string)
  self
end

# File lib/ansi_string.rb, line 184
def dup
  ANSIString.new(@raw.dup)
end

# File lib/ansi_string.rb, line 35
def empty?
  length == 0
end

# File lib/ansi_string.rb, line 27
def insert(position, string)
  if position < 0
    position = @without_ansi.length + position + 1
  end
  self[position...position] = string
  self
end

# File lib/ansi_string.rb, line 211
def inspect
  to_s.inspect
end

# File lib/ansi_string.rb, line 140
def length
  @without_ansi.length
end

# File lib/ansi_string.rb, line 144
def lines
  result = []
  current_string = ""
  @ansi_sequence_locations.map do |location|
    if location[:text] == "\n"
      result << ANSIString.new(current_string + "\n")
      current_string = ""
      next
    end

    location[:text].scan(/.*(?:\n|$)/).each_with_index do |line, i|
      break if line == ""

      if i == 0
        current_string << [
          location[:start_ansi_sequence],
          line,
          location[:end_ansi_sequence]
        ].join
      else
        result << ANSIString.new(current_string)
        current_string = ""
        current_string << [
          location[:start_ansi_sequence],
          line,
          location[:end_ansi_sequence]
        ].join
      end
    end

    if location[:text].end_with?("\n")
      result << ANSIString.new(current_string)
      current_string = ""
      next
    end
  end
  result << ANSIString.new(current_string) if current_string.length > 0
  result
end

# File lib/ansi_string.rb, line 81
def replace(str)
  process_string raw_string_for(str)
  self
end

# File lib/ansi_string.rb, line 86
def reverse
  str = @ansi_sequence_locations.reverse.map do |location|
    [location[:start_ansi_sequence], location[:text].reverse, location[:end_ansi_sequence]].join
  end.join
  ANSIString.new str
end

# File lib/ansi_string.rb, line 77
def rindex(*args)
  @without_ansi.rindex(*args)
end

# File lib/ansi_string.rb, line 93
def scan(pattern)
  results = []
  without_ansi.enum_for(:scan, pattern).each do
    md = Regexp.last_match
    if md.captures.any?
      results << md.captures.map.with_index do |_, i|
        # captures use 1-based indexing
        self[md.begin(i+1)..md.end(i+1)-1]
      end
    else
      results << self[md.begin(0)..md.end(0)-1]
    end
  end
  results
end

# File lib/ansi_string.rb, line 109
def slice(index, length=nil)
  return ANSIString.new("") if length == 0
  range = nil
  index = index.without_ansi if index.is_a?(ANSIString)
  index = Regexp.new Regexp.escape(index) if index.is_a?(String)
  if index.is_a?(Integer)
    length ||= 1
    range = (index..index+length-1)
  elsif index.is_a?(Range)
    range = index
  elsif index.is_a?(Regexp)
    md = @without_ansi.match(index)
    capture_group_index = length || 0
    if md
      capture_group = md.offset(capture_group_index)
      range = (capture_group.first..capture_group.last-1)
    end
  else
    raise(ArgumentError, "Must pass in at least an index or a range.")
  end
  self[range] if range
end

# File lib/ansi_string.rb, line 132
def split(*args)
  raw.split(*args).map { |s| ANSIString.new(s) }
end

# File lib/ansi_string.rb, line 136
def strip
  ANSIString.new raw.strip
end

# File lib/ansi_string.rb, line 188
def sub(pattern, replacement)
  str = ""
  count = 0
  max_count = 1
  index = 0
  @without_ansi.enum_for(:scan, pattern).each do
    md = Regexp.last_match
    str << build_string_with_ansi_for(index...(index + md.begin(0)))
    index = md.end(0)
    break if (count += 1) == max_count
  end
  if index != @without_ansi.length
    str << build_string_with_ansi_for(index..@without_ansi.length)
  end
  nstr = str.gsub /(\033\[[0-9;]*m)(.+?)\033\[0m\1/, '\1\2'
  ANSIString.new(nstr)
end

# File lib/ansi_string.rb, line 206
def to_s
  @raw.dup
end

# File lib/ansi_string.rb, line 341
def build_string_with_ansi_for(range)
  return nil if range.begin > length

  str = ""

  if range.exclude_end?
    range = range.begin..(range.end - 1)
  end

  @ansi_sequence_locations.each do |location|
    # If the given range encompasses part of the location, then we want to
    # include the whole location
    if location[:begins_at] >= range.begin && location[:ends_at] <= range.end
      str << [location[:start_ansi_sequence], location[:text], location[:end_ansi_sequence]].join

    elsif location[:begins_at] >= range.begin && location[:begins_at] <= range.end
      str << [location[:start_ansi_sequence], location[:text][0..(range.end - location[:begins_at])], location[:end_ansi_sequence]].join

    # If the location falls within the given range then  make sure we pull
    # out the bits that we want, and keep ANSI escape sequenece intact while
    # doing so.
    elsif (location[:begins_at] <= range.begin && location[:ends_at] >= range.end) || range.cover?(location[:ends_at])
      start_index = range.begin - location[:begins_at]
      end_index = range.end - location[:begins_at]
      str << [location[:start_ansi_sequence], location[:text][start_index..end_index], location[:end_ansi_sequence]].join
    end
  end
  str
end

# File lib/ansi_string.rb, line 229
def process_string(raw_str)
  @without_ansi = ""
  @ansi_sequence_locations = []
  raw_str.enum_for(:scan, /(\e\[[0-9;]*m)?(.*?)(?=\e\[[0-9;]*m|\Z)/m ).each do
    md = Regexp.last_match
    ansi_sequence, text = md.captures

    previous_sequence_location = @ansi_sequence_locations.last
    if previous_sequence_location
      if ansi_sequence == "\e[0m"
        previous_sequence_location[:end_ansi_sequence] = ansi_sequence
        ansi_sequence = nil
      elsif previous_sequence_location[:start_ansi_sequence] == ansi_sequence
        previous_sequence_location[:text] << text
        previous_sequence_location[:ends_at] += text.length
        previous_sequence_location[:length] += text.length
        @without_ansi << text
        next
      end
    end

    if ansi_sequence.nil? && text.to_s.length == 0
      next
    end

    @ansi_sequence_locations.push(
      begins_at: @without_ansi.length,
      ends_at: [@without_ansi.length + text.length - 1, 0].max,
      length: text.length,
      text: text,
      start_ansi_sequence: ansi_sequence
    )

    @without_ansi << text
  end

  @raw = @ansi_sequence_locations.map do |location|
    [location[:start_ansi_sequence], location[:text], location[:end_ansi_sequence]].compact.join
  end.join

  @ansi_sequence_locations
end

# File lib/ansi_string.rb, line 225
def raw_string_for(str)
  str.is_a?(ANSIString) ? str.raw : str.to_s
end

# File lib/ansi_string.rb, line 272
def replace_in_string(range, replacement_str)
  raise RangeError, "#{range.inspect} out of range" if range.begin > length
  return replacement_str if @ansi_sequence_locations.empty?

  range = range.begin..(range.end - 1) if range.exclude_end?
  str = ""
  @ansi_sequence_locations.each_with_index do |location, j|
    # If the given range encompasses part of the location, then we want to
    # include the whole location
    if location[:begins_at] >= range.begin && location[:ends_at] <= range.end
      end_index = range.end - location[:begins_at] + 1

      str << [
        location[:start_ansi_sequence],
        replacement_str,
        location[:text][end_index..-1],
        location[:end_ansi_sequence]
      ].join

    # If the location falls within the given range then  make sure we pull
    # out the bits that we want, and keep ANSI escape sequenece intact while
    # doing so.
    elsif location[:begins_at] <= range.begin && location[:ends_at] >= range.end
      start_index = range.begin - location[:begins_at]
      end_index = range.end - location[:begins_at] + 1

      str << [
        location[:start_ansi_sequence],
        location[:text][0...start_index],
        replacement_str,
        location[:text][end_index..-1],
        location[:end_ansi_sequence]
      ].join

    elsif location[:ends_at] == range.begin
      start_index = range.begin - location[:begins_at]
      end_index = range.end
      num_chars_to_remove_from_next_location = range.end - location[:ends_at]

      str << [
        location[:start_ansi_sequence],
        location[:text][location[:begins_at]...(location[:begins_at]+start_index)],
        replacement_str,
        location[:text][end_index..-1],
        location[:end_ansi_sequence],
      ].join

      if location=@ansi_sequence_locations[j+1]
        old = location.dup
        location[:text][0...num_chars_to_remove_from_next_location] = ""
        location[:begins_at] += num_chars_to_remove_from_next_location
        location[:ends_at] += num_chars_to_remove_from_next_location
      end

    # If we're pushing onto the end of the string
    elsif range.begin == length && location[:ends_at] == length - 1
      if replacement_str.is_a?(ANSIString)
        str << [location[:start_ansi_sequence], location[:text], location[:end_ansi_sequence], replacement_str].join
      else
        str << [location[:start_ansi_sequence], location[:text], replacement_str, location[:end_ansi_sequence]].join
      end
    else
      str << [location[:start_ansi_sequence], location[:text], location[:end_ansi_sequence]].join
    end
  end

  str
end