# -*- coding: utf-8 -*-

import itertools as it, operator as op, functools as ft
from collections import deque
from contextlib import contextmanager
from io import open
from os.path import join, ismount
import os, re, dbus, fcntl, stat

from . import Collector, Datapoint, user_hz, dev_resolve

import logging
log = logging.getLogger(__name__)


class CGAcct(Collector):


	def __init__(self, *argz, **kwz):
		super(CGAcct, self).__init__(*argz, **kwz)

		self.stuck_list = join(self.conf.cg_root, 'sticky.cgacct')

		# Check which info is available, if any
		self.rc_collectors = list()
		for rc in self.conf.resource_controllers:
			try: rc_collector = getattr(self, rc)
			except AttributeError:
				log.warn( 'Unable to find processor'
					' method for rc {!r} metrics, skipping it'.format(rc) )
				continue
			rc_path = join(self.conf.cg_root, rc)
			if not ismount(rc_path + '/'):
				log.warn(( 'Specified rc path ({}) does not'
					' seem to be a mountpoint, skipping it' ).format(rc_path))
				continue
			log.debug('Using cgacct collector for rc: {}'.format(rc))
			self.rc_collectors.append(rc_collector)

		if not self.rc_collectors: # no point doing anything else
			log.info('No cgroup rcs to poll (rc_collectors), disabling collector')
			self.conf.enabled = False
			return

		# List of cgroup sticky bits, set by this service
		self._stuck_list_file = open(self.stuck_list, 'ab+')
		self._stuck_list = dict()
		fcntl.lockf(self._stuck_list_file, fcntl.LOCK_EX | fcntl.LOCK_NB)
		self._stuck_list_file.seek(0)
		for line in self._stuck_list_file:
			rc, svc = line.strip().split()
			if rc not in self._stuck_list: self._stuck_list[rc] = set()
			self._stuck_list[rc].add(svc)


	def _cg_svc_dir(self, rc, svc=None):
		path = join(self.conf.cg_root, rc, self.conf.systemd_prefix)
		if not svc: return path
		svc = svc.rsplit('@', 1)
		return join(path, '{}.service'.format(svc[0] + '@'), svc[1])\
			if len(svc) > 1 else join(path, '{}.service'.format(svc[0]))

	def _cg_svc_metrics(self, rc, metric, svc_instances):
		return (join( self._cg_svc_dir(rc, svc),
			'{}.{}'.format(rc, metric) ) for svc in svc_instances)

	@contextmanager
	def _cg_metric(self, path, **kwz):
		try:
			with open(path, mode='rb', **kwz) as src: yield src
		except (OSError, IOError) as err:
			log.debug('Failed to open cgroup metric: {}'.format(path, err))
			raise

	@staticmethod
	def _svc_name(svc): return svc.replace('@', '').replace('.', '_')


	@staticmethod
	def _systemd_services():
		for unit in dbus.Interface( dbus.SystemBus().get_object(
					'org.freedesktop.systemd1', '/org/freedesktop/systemd1' ),
				'org.freedesktop.systemd1.Manager' ).ListUnits():
			name, state = it.imap(str, op.itemgetter(0, 4)(unit))
			if name.endswith('.service') and state in ('running', 'start'): yield name[:-8]

	def _systemd_cg_stick(self, rc, services):
		if rc not in self._stuck_list: self._stuck_list[rc] = set()
		stuck_update, stuck = False, set(self._stuck_list[rc])
		services = set(services) # will be filtered and returned
		# Process services, make their cgroups persistent
		for svc in list(services):
			if svc not in stuck:
				svc_tasks = join(self._cg_svc_dir(rc, svc), 'tasks')
				try:
					os.chmod( svc_tasks,
						stat.S_IMODE(os.stat(svc_tasks).st_mode) | stat.S_ISVTX )
				except OSError: services.discard(svc) # not running
				else:
					self._stuck_list[rc].add(svc)
					stuck_update = True
			else: stuck.remove(svc) # to exclude it from the cleanup loop
		# Process stuck cgroups for removed services,
		#  try dropping these, otherwise just unstick and forget
		for svc in stuck:
			svc_dir = self._cg_svc_dir(rc, svc)
			try: os.rmdir(svc_dir)
			except OSError:
				log.debug( 'Non-empty cgroup for'
					' not-running service ({}): {}'.format(svc, svc_dir) )
				svc_tasks = join(svc_dir, 'tasks')
				try:
					os.chmod( svc_tasks,
						stat.S_IMODE(os.stat(svc_tasks).st_mode) & ~stat.S_ISVTX )
				except OSError:
					log.debug('Failed to unstick cgroup tasks file: {}'.format(svc_tasks))
			self._stuck_list[rc].remove(svc)
			stuck_update = True
		# Save list updates, if any
		if stuck_update:
			self._stuck_list_file.seek(0)
			self._stuck_list_file.truncate()
			for rc, stuck in self._stuck_list.viewitems():
				for svc in stuck: self._stuck_list_file.write('{} {}\n'.format(rc, svc))
			self._stuck_list_file.flush()
		return services

	_systemd_sticky_instances = lambda self, rc, services: (
		(self._svc_name(svc), list(svc_instances))
		for svc, svc_instances in it.groupby(
			sorted(set(services).intersection(self._systemd_cg_stick(rc, services))),
			key=lambda k: (k.rsplit('@', 1)[0]+'@' if '@' in k else k) ) )


	def cpuacct( self, services,
			_name = 'processes.services.{}.cpu.{}'.format,
			_stats=('user', 'system') ):
		## "stats" counters (user/system) are reported in USER_HZ - 1/Xth of second
		##  yielded values are in seconds, so counter should have 0-1 range,
		##  when divided by the interval
		## Not parsed: usage (should be sum of percpu)
		for svc, svc_instances in self._systemd_sticky_instances('cpuacct', services):
			if svc == 'total':
				log.warn('Detected service name conflict with "total" aggregation')
				continue
			# user/system jiffies
			stat = dict()
			for path in self._cg_svc_metrics('cpuacct', 'stat', svc_instances):
				try:
					with self._cg_metric(path) as src:
						for name, val in (line.strip().split() for line in src):
							if name not in _stats: continue
							try: stat[name] += int(val)
							except KeyError: stat[name] = int(val)
				except (OSError, IOError): pass
			for name in _stats:
				if name not in stat: continue
				yield Datapoint( _name(svc, name),
					'counter', float(stat[name]) / user_hz, None )
			# usage clicks
			usage = None
			for path in self._cg_svc_metrics('cpuacct', 'usage', svc_instances):
				try:
					with self._cg_metric(path) as src:
						usage = (0 if usage is None else usage) + int(src.read().strip())
				except (OSError, IOError): pass
			if usage is not None:
				yield Datapoint(_name(svc, 'usage'), 'counter', usage, None)


	@staticmethod
	def _iostat(pid, _conv=dict( read_bytes=('r', 1),
			write_bytes=('w', 1), cancelled_write_bytes=('w', -1),
			syscr=('rc', 1), syscw=('wc', 1) )):
		res = dict()
		for line in open('/proc/{}/io'.format(pid), 'rb'):
			line = line.strip()
			if not line: continue
			try: name, val = line.split(':', 1)
			except ValueError:
				log.warn('Unrecognized line format in proc/{}/io: {!r}'.format(pid, line))
				continue
			try: k,m = _conv[name]
			except KeyError: continue
			if k not in res: res[k] = 0
			res[k] += int(val.strip()) * m
		try: res = op.itemgetter('r', 'w', 'rc', 'wc')(res)
		except KeyError:
			raise OSError('Incomplete IO data for pid {}'.format(pid))
		# comm is used to make sure it's the same process
		return open('/proc/{}/comm'.format(pid), 'rb').read(), res

	@staticmethod
	def _read_ids(src):
		return set(it.imap(int, it.ifilter( None,
			it.imap(str.strip, src.readlines()) )))

	def blkio( self, services,
			_caches=deque([dict()], maxlen=2),
			_re_line = re.compile( r'^(?P<dev>\d+:\d+)\s+'
				r'(?P<iotype>Read|Write)\s+(?P<count>\d+)$' ),
			_name = 'processes.services.{}.io.{}'.format ):
		# Caches are for syscall io
		cache_prev = _caches[-1]
		cache_update = dict()

		for svc, svc_instances in self._systemd_sticky_instances('blkio', services):

			## Block IO
			## Only reads/writes are accounted, sync/async is meaningless now,
			##  because only sync ops are counted anyway
			svc_io = dict()
			for metric, src in [ ('bytes', 'io_service_bytes'),
					('time', 'io_service_time'), ('ops', 'io_serviced') ]:
				dst = svc_io.setdefault(metric, dict())
				for path in self._cg_svc_metrics('blkio', src, svc_instances):
					try:
						with self._cg_metric(path) as src:
							for line in src:
								match = _re_line.search(line.strip())
								if not match: continue # "Total" line, empty line
								dev = dev_resolve(*map(int, match.group('dev').split(':')))
								if dev is None: continue
								dev = dst.setdefault(dev, dict())
								iotype, val = match.group('iotype').lower(), int(match.group('count'))
								if iotype not in dev: dev[iotype] = val
								else: dev[iotype] += val
					except (OSError, IOError): pass
			for metric, devs in svc_io.viewitems():
				for dev, vals in devs.viewitems():
					if {'read', 'write'} != frozenset(vals):
						log.warn('Unexpected IO counter types: {}'.format(vals))
						continue
					for k,v in vals.viewitems():
						if not v: continue # no point writing always-zeroes for most devices
						yield Datapoint(_name( svc,
							'blkio.{}.{}_{}'.format(dev, metric, k) ), 'counter', v, None)

			## Syscall IO
			## Counters from blkio seem to be less useful in general,
			##  so /proc/*/io stats are collected for all processes in cgroup
			## Should be very inaccurate if pids are respawning
			tids, pids = set(), set()
			for base in it.imap(ft.partial(
					self._cg_svc_dir, 'blkio' ), svc_instances):
				try:
					with self._cg_metric(join(base, 'tasks')) as src:
						tids.update(self._read_ids(src)) # just to count them
					with self._cg_metric(join(base, 'cgroup.procs')) as src:
						pids.update(self._read_ids(src))
				except (OSError, IOError): continue
			# Process/thread count - only collected here
			yield Datapoint( 'processes.services.'
				'{}.threads'.format(svc), 'gauge', len(tids), None)
			yield Datapoint( 'processes.services.'
				'{}.processes'.format(svc), 'gauge', len(pids), None )

			# Actual io metrics
			svc_update = list()
			for pid in pids:
				try: comm, res = self._iostat(pid)
				except (OSError, IOError): continue
				svc_update.append(((svc, pid, comm), res))
			delta_total = list(it.repeat(0, 4))
			for k,res in svc_update:
				try: delta = map(op.sub, res, cache_prev[k])
				except KeyError: continue
				delta_total = map(op.add, delta, delta_total)
			for k,v in it.izip(['bytes_read', 'bytes_write', 'ops_read', 'ops_write'], delta_total):
				yield Datapoint(_name(svc, k), 'gauge', v, None)
			cache_update.update(svc_update)
		_caches.append(cache_update)


	def memory( self, services,
			_name = 'processes.services.{}.memory.{}'.format ):
		for svc, svc_instances in self._systemd_sticky_instances('memory', services):
			vals = dict()

			for path in self._cg_svc_metrics('memory', 'stat', svc_instances):
				try:
					with self._cg_metric(path) as src:
						for line in src:
							name, val = line.strip().split()
							if not name.startswith('total_'): continue
							name = name[6:]
							val, k = int(val), ( _name(svc, name),
								'gauge' if not name.startswith('pg') else 'counter' )
							if k not in vals: vals[k] = val
							else: vals[k] += val
				except (OSError, IOError): pass

			for prefix in None, 'kmem', 'memsw':
				k = '{}.usage_in_bytes' if prefix else 'usage_in_bytes'
				name = 'usage'
				if prefix: name += '_' + prefix
				for path in self._cg_svc_metrics('memory', k, svc_instances):
					try:
						with self._cg_metric(path) as src:
							vals[_name(svc, name), 'gauge'] = int(src.read().strip())
					except (OSError, IOError): pass

			for (name, val_type), val in vals.viewitems():
				yield Datapoint(name, val_type, val, None)


	def read(self):
		services = list(self._systemd_services())
		for dp in it.chain.from_iterable(
			func(services) for func in self.rc_collectors ): yield dp


collector = CGAcct
