require 'socket'
require 'optparse'

class ShaperQueue
  OVERHEAD = 28                 # IP + UDP
  def initialize(delay, rate)
    @delay = Float(delay)
    @rate = Float(rate)
    @q = []
    @next = Time.now
  end

  def enqueue(buf)
    t = Time.now + @delay
    t = @next if @next > t
    @q.push([t, buf])
    @next = t + (buf.size+OVERHEAD)*8/@rate
  end

  def when
    @q.first.first unless @q.empty?
  end

  def dequeue
    @q.shift[1]
  end
end

class UDPReplySocket < UDPSocket
  attr_accessor :q, :rev
  def initialize(hp, af = Socket::AF_INET)
    super(af)
    bind(*hp)
    @peer = nil
  end
  def urecv
    d, si = recvfrom(65536)
    if @peer != si
      STDOUT.puts "Fixing to #{si[3]}:#{si[1]}"
      @peer = si
    end
    d
  end
  def usend(d)
    raise "Unset peer" unless @peer
    send(d, 0, @peer[3], @peer[1])
  end
end

class UDPBoundSocket < UDPSocket
  attr_accessor :q, :rev
  def initialize(hp, af = Socket::AF_INET)
    super(af)
    connect(*hp)
    @peer = nil
    @hp = hp
  end
  def urecv
    d, si = recvfrom(65536)
    if @peer != si
      STDOUT.puts "Receiving from to #{si[3]}:#{si[1]} -- expected #{@hp[0]}:#{@hp[1]}"
      @peer = si
    end
    d
  end
  def usend(d)
    send(d, 0)
  end
end

if $0 == __FILE__

  OptionParser.new do |opts|
    opts.banner = "Usage: #{$0} options"
    opts.on('-l [HOST:]PORT', String, /(?:(.*):)?(\d+)/,
            "The host/port the initiator sends to") { |value|
      $lhp = value[1, 2]
    }
    opts.on('-c [HOST:]PORT', String, /(?:(.*):)?(\d+)/,
            "The host/port where the respondent is waiting") { |value|
      $chp = value[1, 2]
    }
    opts.on('-b BITRATE', Float, "Bitrate in bit/s (float)") { | $rate |
    }
    opts.on('-d DELAY', Float, "Delay in s (float)") { | $delay | }
    begin
      opts.parse!
      raise OptionParser::ParseError, "All options must be given" unless $lhp and $chp and $rate and $delay
    rescue OptionParser::ParseError => e
      puts "ERROR: #{e.to_s}"
      puts
      puts opts
      exit -1
    end
  end

  so = [UDPReplySocket.new($lhp),
        UDPBoundSocket.new($chp)]
  2.times do |i|
    so[i].q = ShaperQueue.new($delay, $rate)
    so[1-i].rev = so[i]
  end

  loop do
    t = Time.now
    t1 = nil
    so.each do |s|
      s.rev.usend(s.q.dequeue) while s.q.when and s.q.when < t
      if t2 = s.q.when
        t1 = t2 if !t1 or t1 > t2
      end
    end
    t1 -= t if t1
    puts "Delaying for #{t1}"
    r, w, e = select(so, nil, nil, t1)
    r.each do |s|
        s.q.enqueue(s.urecv)
    end if r
  end

end
