#!/usr/bin/env ruby

require 'socket'
require 'timeout'
require 'thread'

class NntpIndexer
	def initialize(host, port='nntp')
		@host, @port = host, port
		@xoversize = 5000
		@release_min_file_count = 3
	end

	def do_test(nthreads, groupname)
		nntpclient = NntpClient.new(@host, @port)
		ans = nntpclient.group(groupname)
		nntpclient.quit

		raise ans if ans !~ /^211 \d+ (\d+) (\d+) /
		min, max = $1.to_i, $2.to_i

		qq = []
		while min < max
			qq << "#{min}-#{min+@xoversize-1}"
			min += @xoversize
		end

		$tot = 0

		puts "#{'max : '.ljust(min.to_s.length+1)}#{max}"
		(1..nthreads).map {
			clt = NntpClient.new(@host, @port)

			if ans = clt.group(groupname) and ans.to_i != 211
				puts ans
				next
			end

			Thread.new(clt, qq) { |c, q|
				while rng = q.shift
					print rng+"\r"
					$stdout.flush
					t = 0
					begin
						c.xover(rng) { |l| t += l.length + 1 }
					rescue Object
						puts $!.class, $!.message
						retry
					end
					Thread.exclusive { $tot += t }
				end
				c.quit
			}
		}.each { |t| t.join if t }
		$tot
	end
end

class NntpClient
	def initialize host, port = 'nntp'
		@host = host
		@port = port
		@curgroup = nil
		@sock = nil
	end

	def debug(*s)
		puts(*s) if $DEBUG
	end
	private :debug

	def reconnect
		@sock.close if @sock
		@sock = TCPSocket.new @host, @port
		debug "< #{@sock.gets}"
		mode 'reader'
		group @curgroup if @curgroup
	end

	def method_missing(m, *a)
		if m == :group
			@curgroup = a.first
		end
		
		cmd = ([m] + a).join ' '

		retried = false
		begin
			debug "> #{cmd}"
			@sock.puts cmd
			
			str = Timeout::timeout(20, RuntimeError) { @sock.gets }.chomp! or raise 'partial read'
			debug "< #{str}"
		rescue
			raise if retried or m == :quit
			retried = true

			reconnect
			retry
		end

		# no yield if response != 2xx
		if block_given? and str[0] == ?2
			begin
				sstr = nil
				loop do
					sstr = @sock.gets.chomp! or raise 'partial read'
					break if sstr == '.'
					sstr = sstr[1..-1] if sstr[0] == ?. and sstr[1] == ?.
					yield sstr
				end
			ensure
				# consume data even on 'break'
				Timeout::timeout(30) {
					sstr = @sock.gets.chomp while sstr != '.'
				}
			end
		end

		str
	end
end

if $0 == __FILE__
n = NntpIndexer.new('news.free.fr')
group = ARGV.shift
5.times { |i|
	tstart = Time.now
	nbytes = n.do_test i, group
	delay = Time.now - tstart
	puts "with #{i} threads : #{delay} seconds, #{nbytes} bytes => #{(nbytes/delay/1024).to_i} ko/s"
}
end