require 'socket'
require 'sqlite3'

class NntpClient
	def initialize host
		@host = host
		@curgroup = nil
		reconnect
	end

	def reconnect
		@sock = TCPSocket.new @host, 'nntp'
		puts ' [+] connected'
		s = @sock.gets
		puts s if $VERBOSE
		group @curgroup if @curgroup
	end

	def method_missing(m, *a)
		if m == :group
			@curgroup = a.first
		end
		
		cmd = m.to_s << a.map { |s| ' ' << s.to_s }.join

		puts "> #{cmd}" if $VERBOSE
		begin
			@sock.puts cmd
		rescue
			return if m == :quit
			reconnect
			@sock.puts cmd
		end

		str = @sock.gets.chomp! or raise 'partial read'
		puts str if $VERBOSE

		# no yield if response != 2xx
		if block_given? and str[0] == ?2
			loop do
				str = @sock.gets.chomp! or raise 'partial read'
				return if str == '.'
				str = str[1..-1] if str[0] == ?.
				yield str
			end
		end
		str
	end
end

class NntpDB
	attr_reader :db
	def initialize name
		@created = true if not File.exist? name

		@db = SQLite3::Database.new name
		if @created
			@db.execute_batch <<EOSQL
create table groups (
 id integer primary key autoincrement,
 name varchar(255),
 low integer,
 high integer
);
create unique index name on groups (name);

create table articles (
 groupid integer,
 id integer,
 subject varchar(1024),
 poster varchar(255),
 date varchar(255),
 msgid varchar(255),
 bytes integer,
 lines integer,
 primary key (groupid, id)
);
create index groupid on articles (groupid);
create index msgid on articles (msgid);
create index subject on articles (subject);
EOSQL
		end
	end

	def update_group_list nntpclient
		@db.transaction
		if @created
			stmt = @db.prepare 'insert into groups (name, low, high) values (?, ?, ?)'
			nntpclient.list { |g|
				name, high, low = g.split ' '
				high, low = high.to_i, low.to_i
				next unless yield name, low, high if block_given?
				stmt.execute name, low, high
			}
		else
		        # XXX delete obsoleted groups
			stmt = @db.prepare 'update groups set low=?, high=? where id=?'
			nntpclient.list { |g|
				name, high, low = g.split ' '
				high, low = high.to_i, low.to_i
				next unless yield name, low, high if block_given?
				if id = @db.get_first_value('select id from groups where name=?', name)
					stmt.execute low, high, id
				else
					@db.execute 'insert into groups (name, low, high) values (?, ?, ?)', name, low, high
				end
			}
		end
		stmt.close
		@db.commit
	end

	def method_missing(*a, &b)
		@db.send *a, &b
	end
end

def update_news(server, file, min, max)

db = NntpDB.new file

highest_seen = {}
db.execute('select groupid, max(id) from articles group by groupid') { |gid, mx|
	highest_seen[gid.to_i] = mx.to_i
}

s = NntpClient.new server

begin

db.update_group_list(s) { |name, low, high|
	if block_given?
		yield name
	else
		count = high-low
		count >= min and count < max
	end
}

grouplist = db.execute 'select id, name from groups'

puts " [+] number of groups : #{grouplist.length}"

step = 100000


grpcnt = 0
grouplist.sort_by { |id, name| name }.each { |groupid, groupname|
	next if ARGV[0] and groupname !~ /#{ARGV[0]}/

	groupid = groupid.to_i

	puts " [+] group #{grpcnt += 1}/#{grouplist.length} : #{groupname}"

	low, high = db.get_first_row('select low, high from groups where id=?', groupid).map { |i| i.to_i }
	next if highest_seen[groupid] == high

	ans = s.group(groupname)
	resp, count, low, high = ans.split(' ').map { |i| i.to_i }

	if resp != 211
		puts " [-] response #{ans} - deleting"
		db.transaction {
			db.execute 'delete from articles where groupid=?', groupid
			db.execute 'delete from groups where id=?', groupid
		}
		next
	end


	db.transaction {
		db.execute 'delete from articles where groupid=? and id<?', groupid, low
		db.execute 'update groups set low=? where id=?', low, groupid
	} if db.get_first_value('select * from articles where groupid=? and id<? limit 1', groupid, low)
	db.execute('update groups set high=? where id=?', high, groupid) if db.get_first_value('select high from groups where id=?', groupid).to_i < high
	
	next if low > high
	
	# nil for unchecked groups
	low = (highest_seen[groupid] or low-1).to_i + 1
	while high > low
		nextlow = low + step
		nextlow = high if nextlow > high - step/3
		retried = 0
		puts " [+] xover #{low}-#{nextlow}#{" (#{(nextlow-min)*100/(high-min)}%)" if nextlow != high}" # vim sux "
		begin
			db.transaction
			stmt = db.prepare 'insert into articles (groupid, id, subject, poster, date, msgid, bytes, lines) values (?, ?, ?, ?, ?, ?, ?, ?)'
			s.xover("#{low}-#{nextlow}") { |e|
				id, subject, from, date, mid, ref, bytes, lines = e.split "\t"
				stmt.execute groupid, id.to_i, subject, from, date, mid, bytes.to_i, lines.to_i
			}
			db.execute 'update groups set high=? where id=?', nextlow, groupid
			low = nextlow + 1
		rescue
			stmt.close
			db.rollback
			puts " [-] Exception: #{$!.class} : #{$!.message rescue nil}"
			raise if retried > 4
			step /= 2
			step = 5000 if step < 5000
			retried += 1
			s.reconnect
			retry
		rescue Object
			stmt.close
			db.rollback
			raise
		else
			stmt.close
			db.commit
			step *= 2 if retried == 0 and step < 100000 and nextlow != high
		end
	end
}

puts "total number of articles indexed : #{db.get_first_value 'select count(*) from articles'}"

ensure
	db.close
	s.quit rescue nil
end

end