require 'socket'
require 'timeout'

class StompClient
  def initialize(host, port)
    @sock = TCPSocket.new(host, port)
    @buf = +''
  end

  def send_frame(command, headers = {}, body = '')
    frame = command + "\n"
    headers.each { |k, v| frame << "#{k}:#{v}\n" }
    frame << "\n"
    frame << body
    frame << "\0"
    @sock.write(frame)
  end

  def read_bytes(n)
    while @buf.bytesize < n
      chunk = @sock.readpartial(4096)
      raise 'connection closed' if chunk.nil? || chunk.empty?
      @buf << chunk
    end
    @buf.slice!(0, n)
  end

  def read_frame(timeout: 5)
    Timeout.timeout(timeout) do
      until @buf.include?("\n\n")
        chunk = @sock.readpartial(4096)
        raise 'connection closed' if chunk.nil? || chunk.empty?
        @buf << chunk
      end
      header_end = @buf.index("\n\n") + 2
      header_blob = @buf.slice!(0, header_end)
      lines = header_blob.strip.split("\n")
      command = lines.shift
      headers = {}
      lines.each do |line|
        key, value = line.split(':', 2)
        headers[key] = value
      end
      if headers['content-length']
        body = read_bytes(headers['content-length'].to_i)
        read_bytes(1) if @buf.bytesize.positive? && @buf.getbyte(0) == 0
      else
        until @buf.index("\0")
          chunk = @sock.readpartial(4096)
          raise 'connection closed' if chunk.nil? || chunk.empty?
          @buf << chunk
        end
        body = @buf.slice!(0, @buf.index("\0"))
      end
      [command, headers, body]
    end
  end

  def close
    @sock.close
  end
end