#!/usr/bin/env python3
'''Client for the routerstats collection'''

import socket
import logging
import time
import os
import sys
import argparse
import configparser

import rrdtool

logging.basicConfig(
    format='%(asctime)s %(funcName)20s %(levelname)-8s %(message)s',
    level=logging.INFO,
    datefmt='%Y-%m-%d %H:%M:%S')

class NotConnected(Exception):
	'''Raise to make it known that you are not connected'''

class RouterstatsClient():
	'''Client itself'''
	def __init__(self, host, port, passwd_file):
		self.host = host
		self.port = port
		self.connected = False
		self.sock = None
		self.passwd_file = passwd_file
		self.received_lines = 0

	def connect(self):
		'''Do connect'''
		self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		try:
			self.sock.connect((self.host, self.port))
			self.sock.settimeout(1)
			logging.info('Connected to %s', self.host)
			self.login()
			self.connected = True
			self.received_lines = 0
		except ConnectionRefusedError as error:
			logging.error('Could not connect to %s:%s', self.host, self.port)
			raise ConnectionRefusedError from error
		except ConnectionError as error:
			logging.error('Could not connect to %s:%s: %s',
									self.host,
									self.port,
									error)
		except TimeoutError as error:
			logging.error('Timed out')

	def login(self):
		'''Do the login routine.
			Wait for server to greet us with "hello"
			reply with the password
			wait for "Welcome"
		'''

		logging.debug('Logging in')
		try:
			hello = self.sock.recv(5)
		except TimeoutError as exception:
			logging.error('Timeout while waiting for Hello')
			raise ConnectionError('Timed out waiting for Hello when connecting') from exception
		if not hello == b'Hello':
			logging.error('No Hello from server: %s', hello)
			raise ConnectionError('Server did not greet us with Hello during login')
		with open(self.passwd_file, 'r', encoding='utf-8') as passwd_file:
			passwd = passwd_file.readline()
			passwd = passwd.rstrip()
		logging.debug('Sending password: %s', passwd)
		self.sock.send(passwd.encode('utf-8'))
		try:
			response = self.sock.recv(7)
			if not response == b'Welcome':
				logging.error('Not Welcome: %s', response)
				raise ConnectionError('We are not greeted with Welcome after sending password')
			return True
		except TimeoutError as exception:
			raise ConnectionError(
						'Timed out while waiting for server to greet us after sending password') from exception

	def send(self, tosend):
		'''Send some data'''
		if self.connected:
			logging.debug('Sending %s', tosend)
			send_bytes = bytes(tosend + "\n", 'utf-8')
			try:
				self.sock.send(send_bytes)
			except OSError as error:
				logging.error('Cound not send to server: %s', error)
				self.connected = False
				raise NotConnected from error
			try:
				self.sock.getpeername()
			except OSError as error:
				logging.error('Could not send to server: %s', error)
				self.connected = False
				raise NotConnected from error
		else:
			logging.error('Not connected to server')
			raise NotConnected('Not connected to server')

	def recv(self):
		'''Receive some data'''
		if self.connected is not True:
			logging.error('Trying to recv when not connected')
			raise NotConnected
		line = b''
		blanks = 0
		while True:
			try:
				toret = self.sock.recv(1)
				if toret:
					line += toret
					if blanks > 0:
						blanks -= 1
					if line.endswith(b'\n'):
						#We're done for now, returning value
						line = line.strip().decode('utf-8')
						logging.debug('Received from server: %s', line)
						self.received_lines += 1
						if not self.received_lines % 1000:
							logging.info('Received %s lines', self.received_lines)
						return line
				else:
					blanks += 1
					if blanks >= 5:
						blanks = 0
#						break
						logging.debug('Too many blank reads, and still no complete line')
						self.connected = False
						raise NotConnected
			except TimeoutError:
				if line == b'':
					break
				logging.error('Timeout while fetching data, got: %s' ,line)
				break
			except OSError as error:
				if str(error) == 'timed out':
					#This is expected, as it is how we loop
					break
				logging.error('OSError: %s', error)
				logging.debug('Got: %s', line)
				self.connected = False
				raise NotConnected(error) from error
		return None

def handle_server_msg(received: str):
	'''Do something about something from the server'''
	try:
		timestamp, zone = received.split()
	except ValueError:
		logging.error('Could not parse %s into two distinct values', received)
		return None
	try:
		timestamp = int(timestamp)
	except ValueError:
		logging.error('Could not parse %s as an int', timestamp)
		return None
	toret = {'timestamp': timestamp, 'net_dnat': 0, 'loc-net': 0}
	try:
		toret[zone] += 1
	except KeyError as error:
		logging.debug('Ignoring zone: %s', error)
	logging.debug('Parsed to: %s', toret)
	return toret

def handle(received, client):
	'''Do things with whatever came from the server'''
	if not received:
		return None
	if received == 'ping':
		client.send('pong')
		return True
	if received is not None:
		if received:
			return handle_server_msg(received)
	return None

class UpdateRRD:
	'''Handle updates to the rrd-file, since we cannot update more than once every second'''
	def __init__(self, rrdfile):
		self.rrdfile = rrdfile
		if os.path.isfile(rrdfile) is not True:
			self.create()
		self.toupdate = {'timestamp': None, 'net_dnat': 0, 'loc-net': 0}
		self.freshdict = self.toupdate.copy()

	def create(self):
		'''Create self.rrdfile'''
		rrdtool.create(
								self.rrdfile,
								"--start",
								"1000000000",
								"--step",
								"10",
								"DS:net_dnat:ABSOLUTE:10:0:U",
								"DS:loc-net:ABSOLUTE:10:0:U",
								"RRA:AVERAGE:0.5:1:1d",
								"RRA:AVERAGE:0.5:1:1M")
		logging.debug('Created rrdfile %s', self.rrdfile)

	def push(self):
		'''Send data to rrdtool'''
		#Make sure we're not doing something stupid..
		info = rrdtool.info(self.rrdfile)
		if self.toupdate['timestamp'] is None:
			return False
		if info['last_update'] > self.toupdate['timestamp']:
			logging.error(
				'Trying to update when rrdfile is newer than our timestamp. Ignoring line and resetting.')
			self.toupdate = self.freshdict.copy()
			return False
		if info['last_update'] == self.toupdate['timestamp']:
			logging.error('last update and toupdate timestamp are the same, this should not happen')
			self.toupdate = self.freshdict.copy()
			return False
		try:
			rrdtool.update(
						self.rrdfile,
						str(self.toupdate['timestamp'])
						+ ':'
						+ str(self.toupdate['net_dnat'])
						+ ':'
						+ str(self.toupdate['loc-net']))
			self.toupdate = self.freshdict.copy()
			logging.debug('Updated rrdfile')
			return True
		except rrdtool.OperationalError as error:
			if str(error) == 'could not lock RRD':
				return False
			logging.error(str(error))
			return False

	def add(self, input_dict):
		'''Add data to be added later'''
		if self.toupdate['timestamp'] is None:
			#Never touched, just overwrite with whatever we got
			self.toupdate = input_dict
		elif input_dict['timestamp'] > self.toupdate['timestamp']:
			#What we get is fresher than what we have, and noone has done a push yet
			self.push()
			self.toupdate = input_dict
		elif input_dict['timestamp'] == self.toupdate['timestamp']:
			#Same timestamp, just up values and be happy
			self.toupdate['net_dnat'] += input_dict['net_dnat']
			self.toupdate['loc-net'] += input_dict['loc-net']
		elif input_dict['timestamp'] < self.toupdate['timestamp']:
			diff = self.toupdate['timestamp'] - input_dict['timestamp']
			if diff <= 5:
				#Might be because something is a bit slow sometimes?
				input_dict['timestamp'] = self.toupdate['timestamp']
			else:
				logging.error('Newly fetched data is older than what we have in the queue already. Passing.')
		else:
			logging.error('Not sure what to do here? %s: %s', input_dict, self.toupdate)

def initial_connect(client):
	'''Loop here until first connection is made'''
	while True:
		try:
			client.connect()
			break
		except ConnectionRefusedError:
			time.sleep(1)

def loop(client, rrdupdater):
	'''Main loop'''
	initial_connect(client)

	tries = 0
	loops = 0
	while True:
		try:
			retval = handle(client.recv(), client)
			if retval:
				loops = 0
				if retval is True:
					pass
				else:
					rrdupdater.add(retval)
			else:
				loops += 1
				if loops >= 5:
					#Want to wait until at least 5 secs have passed since last actual data fetch..
					rrdupdater.push()
				if loops >= 60:
					logging.error(
						'No data in 60 seconds. We expect a ping/pong every 30. Lost connection, probably')
					loops = 0
					raise NotConnected
		except NotConnected:
			try:
				client.connect()
				tries = 0
			except ConnectionRefusedError as error:
				logging.debug('%s', error)
				tries += 1
				tries = min(tries, 5)
				time.sleep(tries)
		except ConnectionRefusedError as error:
			logging.debug('%s', error)
			time.sleep(1)
		except KeyboardInterrupt:
			break

def main():
	'''Main is main'''
	config_section = 'client'

	config = configparser.ConfigParser()
	parser = argparse.ArgumentParser(exit_on_error=False, add_help=False)

	parser.add_argument('-c', '--config', help='config file to load')
	parser.add_argument('-d', '--debug', action='store_true', help='enable debug')

	args, _ = parser.parse_known_args()

	if args.debug:
		logging.root.setLevel(logging.DEBUG)

	logging.debug('Starting as PID %s', os.getpid())

	found = False
	config_dirs = ('/etc/routerstats/', '/usr/local/etc/routerstats/', '/opt/routerstats/', './')
	if args.config:
		if os.path.isfile(args.config):
			config.read(args.config)
			found = True
		else:
			logging.error('Specified config file does not exist: %s', args.config)
	else:
		logging.debug('Trying to find config')
		#Try to find in "usual" places
		for directory in config_dirs:
			trytoread = directory + 'routerstats.config'
			if os.path.isfile(trytoread):
				logging.debug('Reading config file %s', trytoread)
				config.read(trytoread)
				found = True
	if not found:
		logging.error('routerstats.config not found in %s', config_dirs)
		sys.exit(0)

	parser.add_argument(
					'-r',
					'--host',
					dest='host',
					help='ip/hostname to connect to',
					default=config[config_section]['host'])
	parser.add_argument(
					'-p',
					'--port',
					dest='port',
					type=int,
					help='port to connect to',
					default=config[config_section]['port'])
	parser.add_argument(
					'-v',
					'--vardir',
					dest='vardir',
					help='directory for storing rrd file',
					default=config[config_section]['var_dir'])
	parser.add_argument(
					'-w',
					'--pwdfile',
					dest='passwd_file',
					help='password file',
					default=config[config_section]['passwd_file'])
	parser.add_argument(
					'-h',
					'--help',
					help='show this help and exit',
					action='store_true',
					default=False)

	args = parser.parse_args()
	logging.debug(args)

	if args.help:
		parser.print_help()
		sys.exit()

	#Make sure the file specified is to be found..
	if not os.path.isfile(args.passwd_file):
		logging.error('Cannot find passwd-file %s', args.passwd_file)
		sys.exit()

	if not os.path.isdir(args.vardir):
		logging.error('Cannot find var dir %s', args.vardir)

	rrdupdater = UpdateRRD(args.vardir + '/routerstats.rrd')
	client = RouterstatsClient(args.host, args.port, args.passwd_file)
	loop(client, rrdupdater)

if __name__ == '__main__':
	main()