class Segregate

# File lib/segregate.rb, line 31
def initialize callback = nil, *args, **kwargs
  @debug = kwargs[:debug] ? true : false
  @callback = callback
  @http_version = [nil, nil]

  # :request, :response
  @type = Juncture.new :request, :response
  @state = Juncture.new :waiting, :headers, :body, :done, default: :waiting

  @headers = Hashie::Mash.new
  @body = ""

  @stashed_data = ""
  @stashed_body = ""

  @header_order = []
end

# File lib/segregate.rb, line 126
def build_body raw_message
  new_body, size = deflate_if_needed(@body)
  if @headers['content-length']
    raw_message << new_body
    raw_message << "\r\n\r\n"
  elsif @headers['transfer-encoding'] == 'chunked'
    raw_message << "%s\r\n" % (size.to_s(16))
    raw_message << new_body + "\r\n"
    raw_message << "0\r\n\r\n"
  end
end

# File lib/segregate.rb, line 111
def build_headers raw_message
  request? ? raw_message << request_line + "\r\n" : raw_message << status_line + "\r\n"
  @header_order.each do |header|
    values = @headers[header.downcase]
    if values.kind_of?(Array)
      values.each do |value|
        raw_message << "%s: %s\r\n" % [header, value]
      end
    else
      raw_message << "%s: %s\r\n" % [header, values]
    end
  end
  raw_message << "\r\n"
end

# File lib/segregate.rb, line 25
def debug message
  if @debug
    puts "DEBUG: " + message.to_s
  end
end

# File lib/segregate.rb, line 138
def deflate_if_needed(body)
  return [body, body.size] unless gzip_encoded_body?

  str_io = StringIO.new('w')
  w_gz = Zlib::GzipWriter.new(str_io)
  w_gz.write(body)
  w_gz.close

  # To specify the number of bytes for a gzip string
  # we should check for the buffer size instead of the string size
  [str_io.string, str_io.size]
end

# File lib/segregate.rb, line 61
def done?
  @state >= :done
end

# File lib/segregate.rb, line 57
def headers_complete?
  @state > :headers
end

# File lib/segregate.rb, line 77
def major_http_version
  http_version[0]
end

# File lib/segregate.rb, line 81
def major_http_version= value
  http_version[0] = value
end

# File lib/segregate.rb, line 13
def method_missing meth, *args, &block
  if @uri.respond_to? meth
    @uri.send meth, *args, &block
  else
    super
  end
end

# File lib/segregate.rb, line 85
def minor_http_version
  http_version[1]
end

# File lib/segregate.rb, line 89
def minor_http_version= value
  http_version[1] = value
end

# File lib/segregate.rb, line 151
def parse_data data
  data = StringIO.new(@stashed_data + data)
  @stashed_data = ""

  while !data.eof? && @state < :done
    line, complete_line = get_next_line data
    complete_line ? parse_line(line) : @stashed_data = line
  end

  data.close
end

# File lib/segregate.rb, line 163
def parse_line line
  case @state.state
  when :waiting
    read_in_first_line line
  when :headers
    read_in_headers line
  when :body
    read_in_body line
  end

  @callback.on_message_complete self if @callback.respond_to?(:on_message_complete) && done?
end

# File lib/segregate.rb, line 101
def raw_data
  raw_message = ""
  update_content_length

  build_headers raw_message
  build_body raw_message

  return raw_message
end

# File lib/segregate.rb, line 49
def request?
  @type == :request
end

# File lib/segregate.rb, line 65
def request_line
  request? ? "%s %s HTTP/%d.%d" % [request_method, request_url.to_s, *http_version] : nil
end

# File lib/segregate.rb, line 73
def request_url
  uri ? uri.to_s : nil
end

# File lib/segregate.rb, line 21
def respond_to?(meth, include_private = false)
  @uri.respond_to?(meth, include_private) || super
end

# File lib/segregate.rb, line 53
def response?
  @type == :response
end

# File lib/segregate.rb, line 69
def status_line
  response? ? "HTTP/%d.%d %d %s" % [*http_version, status_code, status_phrase] : nil
end

# File lib/segregate.rb, line 93
def update_content_length
  unless @body.empty?
    if @headers['content-length']
      @headers['content-length'] = @body.length.to_s
    end
  end
end

# File lib/segregate.rb, line 178
def get_next_line data
  if @headers['content-length'] && @state >= :body
    line = data.readline("\r\n")
    @inital_line = line
    [line, true]
  else
    line = data.readline("\r\n")
    @inital_line = line
    result = line.end_with?("\r\n")
    line = line[0..-3] if result
    [line, result]
  end
end

# File lib/segregate.rb, line 272
def gzip_encoded_body?
  headers.key?('content-encoding') && headers['content-encoding'].eql?('gzip')
end

# File lib/segregate.rb, line 264
def inflate_body_if_needed(body)
  return body unless gzip_encoded_body?

  gzip_str  = StringIO.new(body)
  gzip_body = Zlib::GzipReader.new( gzip_str )
  gzip_body.read()
end

# File lib/segregate.rb, line 249
def parse_body line
  line = @stashed_body + line
  @stashed_body = ""

  if line.length >= headers['content-length'].to_i
    # Removing delimiters characters from the HTTP protocol
    line = line[0..-3]
    @body = inflate_body_if_needed(line)
    @callback.on_body @body if @callback.respond_to?(:on_body)
    @state.next
  else
    @stashed_body = line
  end
end

# File lib/segregate.rb, line 276
def parse_chunked_data line
  @chunked_body_state ||= Juncture.new :size, :chunk, default: :size

  case @chunked_body_state.state
  when :size
    parse_chunked_data_size line.to_i(16)
  when :chunk
    parse_chunked_data_chunk line
  end
end

# File lib/segregate.rb, line 296
def parse_chunked_data_chunk line
  line = @stashed_body + line
  @stashed_body = ""

  if line.length >= @chunk_size
    @body << line
    @original_body = @body
    @body = inflate_body_if_needed(@body)
    @callback.on_body line if @callback.respond_to?(:on_body)
    @chunked_body_state.next
  else
    @stashed_body = @inital_line.end_with?("\r\n") ? (line + "\r\n") : line
  end
end

# File lib/segregate.rb, line 287
def parse_chunked_data_size chunk_size
  if chunk_size == 0
    @state.next
  else
    @chunk_size = chunk_size
  end
  @chunked_body_state.next
end

# File lib/segregate.rb, line 206
def parse_request_line line
  @type.set :request
  @request_method, url, @http_version[0], @http_version[1] = line.scan(REQUEST_LINE).flatten
  @http_version.map! {|v| v.to_i}
  @uri = URI.parse url

  @callback.on_request_line self if @callback.respond_to?(:on_request_line)
end

# File lib/segregate.rb, line 215
def parse_status_line line
  @type.set :response
  @http_version[0], @http_version[1], code, @status_phrase = line.scan(STATUS_LINE).flatten
  @http_version.map! {|v| v.to_i}
  @status_code = code.to_i

  @callback.on_status_line self if @callback.respond_to?(:on_status_line)
end

# File lib/segregate.rb, line 241
def read_in_body line
  if headers['content-length']
    parse_body line
  elsif headers['transfer-encoding'] == 'chunked'
    parse_chunked_data line
  end
end

# File lib/segregate.rb, line 192
def read_in_first_line line
  @callback.on_message_begin self if @callback.respond_to?(:on_message_begin)

  if line =~ REQUEST_LINE
    parse_request_line line
  elsif line =~ STATUS_LINE
    parse_status_line line
  else
    raise "ERROR: Unknown first line: %s" % line
  end

  @state.next
end

# File lib/segregate.rb, line 224
def read_in_headers line
  if line.empty?
    @state.next
  else
    key, value = line.split(": ",2)
    @header_order << key unless @header_order.include?(key)
    save_value(@headers, key, value)
  end

  if headers_complete?
    @callback.on_headers_complete self if @callback.respond_to?(:on_headers_complete)
    if headers['transfer-encoding'].nil? && (headers['content-length'] == "0" || headers['content-length'].nil?)
      @state.set :done
    end
  end
end

# File lib/segregate.rb, line 311
def save_value(store_hash, key, value)
  key = key.downcase
  if store_hash.key?(key)
    store_hash[key] = [] << store_hash[key] unless store_hash[key].kind_of?(Array)
    store_hash[key] << value
  else
    store_hash[key] = value
  end
end