"""
Copyright 2016 Deepgram

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import os
import json
import signal
import atexit
import time
import sys
import argparse
import logging

from . import __version__, __homepage__
from .utils import logcolor
from . import Kurfile
from .plugins import Plugin
from .engine import JinjaEngine

logger = logging.getLogger(__name__)

###############################################################################
def parse_kurfile(filename, engine, parse=True):
	""" Parses a Kurfile.

		# Arguments

		filename: str. The path to the Kurfile to load.

		# Return value

		Kurfile instance
	"""
	spec = Kurfile(filename, engine)
	if parse:
		spec.parse()
	return spec

###############################################################################
def dump(args):
	""" Dumps the Kurfile to stdout as a JSON blob.
	"""
	spec = parse_kurfile(args.kurfile, args.engine, parse=not args.pre_parse)
	print(json.dumps(spec.data, sort_keys=True, indent=4))

###############################################################################
def train(args):
	""" Trains a model.
	"""
	spec = parse_kurfile(args.kurfile, args.engine)
	func = spec.get_training_function()
	func(step=args.step)

###############################################################################
def test(args):
	""" Tests a model.
	"""
	spec = parse_kurfile(args.kurfile, args.engine)
	func = spec.get_testing_function()
	func(step=args.step)

###############################################################################
def evaluate(args):
	""" Evaluates a model.
	"""
	spec = parse_kurfile(args.kurfile, args.engine)
	func = spec.get_evaluation_function()
	func(step=args.step)

###############################################################################
def build(args):
	""" Builds a model.
	"""
	spec = parse_kurfile(args.kurfile, args.engine)

	if args.compile == 'auto':
		result = []
		for section in ('train', 'test', 'evaluate'):
			if section in spec.data:
				result.append((section, 'data' in spec.data[section]))
		if not result:
			logger.info('Trying to build a bare model.')
			args.compile = 'none'
		else:
			args.compile, has_data = sorted(result, key=lambda x: not x[1])[0]
			logger.info('Trying to build a "%s" model.', args.compile)
			if not has_data:
				logger.warning('There is not data defined for this model, '
					'so we will be running as if --bare was specified.')
	elif args.compile == 'none':
		logger.info('Trying to build a bare model.')
	else:
		logger.info('Trying to build a "%s" model.', args.compile)

	if args.bare or args.compile == 'none':
		provider = None
	else:
		providers = spec.get_provider(
			args.compile,
			accept_many=args.compile == 'test'
		)
		provider = Kurfile.find_default_provider(providers)

	spec.get_model(provider)

	if args.compile == 'none':
		return
	elif args.compile == 'train':
		target = spec.get_trainer(with_optimizer=True)
	elif args.compile == 'test':
		target = spec.get_trainer(with_optimizer=False)
	elif args.compile == 'evaluate':
		target = spec.get_evaluator()
	else:
		logger.error('Unhandled compilation target: %s. This is a bug.',
			args.compile)
		return 1

	target.compile()

###############################################################################
def prepare_data(args):
	""" Prepares a model's data provider.
	"""
	spec = parse_kurfile(args.kurfile, args.engine)

	if args.target == 'auto':
		result = None
		for section in ('train', 'validate', 'test', 'evaluate'):
			if section in spec.data and 'data' in spec.data[section]:
				result = section
				break
		if result is None:
			raise ValueError('No data sections were found in the Kurfile.')
		args.target = result

	logger.info('Preparing data sources for: %s', args.target)

	providers = spec.get_provider(
		args.target,
		accept_many=args.target == 'test'
	)

	if args.assemble:

		default_provider = Kurfile.find_default_provider(providers)
		spec.get_model(default_provider)

		if args.target == 'train':
			target = spec.get_trainer(with_optimizer=True)
		elif args.target == 'test':
			target = spec.get_trainer(with_optimizer=False)
		elif args.target == 'evaluate':
			target = spec.get_evaluator()
		else:
			logger.error('Unhandled assembly target: %s. This is a bug.',
				args.target)
			return 1

		target.compile(assemble_only=True)

	for k, provider in providers.items():
		if len(providers) > 1:
			print('Provider:', k)

		batch = None
		for batch in provider:
			break
		if batch is None:
			logger.error('No batches were produced.')
			continue

		num_entries = None
		keys = sorted(batch.keys())
		num_entries = len(batch[keys[0]])
		for entry in range(num_entries):
			if args.number is None or entry < args.number or \
				(entry - num_entries >= args.number):
				print('Entry {}/{}:'.format(entry+1, num_entries))
				for key in keys:
					print('  {}: {}'.format(key, batch[key][entry]))

		if num_entries is None:
			logger.error('No data sources was produced.')
			continue

###############################################################################
def version(args):							# pylint: disable=unused-argument
	""" Prints the Kur version and exits.
	"""
	print('Kur, by Deepgram -- deep learning made easy')
	print('Version: {}'.format(__version__))
	print('Homepage: {}'.format(__homepage__))

###############################################################################
def do_monitor(args):
	""" Handle "monitor" mode.
	"""

	# If we aren't running in monitor mode, then we are done.
	if not args.monitor:
		return

	# This is the main retry loop.
	while True:
		# Fork the process.
		logger.info('Forking child process.')
		pid = os.fork()

		# If we are the child, leave this function and work.
		if pid == 0:
			logger.debug('We are a newly spawned child process.')
			return

		logger.debug('Child process spawned: %d', pid)

		# Wait for the child to die. If we die first, kill the child.
		atexit.register(kill_process, pid)
		try:
			_, exit_status = os.waitpid(pid, 0)
		except KeyboardInterrupt:
			break
		atexit.unregister(kill_process)

		# Process the exit code.
		signal_number = exit_status & 0xFF
		exit_code = (exit_status >> 8) & 0xFF
		core_dump = bool(0x80 & signal_number)

		if signal_number == 0:
			logger.info('Child process exited with exit code: %d.', exit_code)
		else:
			logger.info('Child process exited with signal %d (core dump: %s).',
				signal_number, core_dump)

		retry = False
		if os.WIFSIGNALED(exit_status):
			if os.WTERMSIG(exit_status) == signal.SIGSEGV:
				logger.error('Child process seg faulted.')
				retry = True

		if not retry:
			break

	sys.exit(0)

###############################################################################
def kill_process(pid):
	""" Kills a child process by PID.
	"""

	# Maximum time we wait (in seconds) before we send SIGKILL.
	max_timeout = 60

	# Terminate child process
	logger.debug('Sending Ctrl+C to the child process %d', pid)
	os.kill(pid, signal.SIGINT)

	start = time.time()
	while True:
		now = time.time()

		# Check the result.
		result = os.waitpid(pid, os.WNOHANG)
		if result != (0, 0):
			# The child process is dead.
			break

		# Check the timeout.
		if now - start > max_timeout:
			# We've waited too long.
			os.kill(pid, signal.SIGKILL)
			break

		# Keep waiting.
		logger.debug('Waiting patiently...')
		time.sleep(0.5)

###############################################################################
def load_plugins(plugin_dir):
	""" Loads the Kur plugins.
	"""
	if plugin_dir:
		if plugin_dir.lower() == 'none':
			return
		Plugin.PLUGIN_DIR = plugin_dir

	for plugin in Plugin.get_enabled_plugins():
		try:
			plugin.load()
		except:
			logger.error('Failed to load plugin "%s". It may not be installed '
				'correctly, or it may contains errors.', plugin)
			raise

###############################################################################
def install_plugin(args):
	""" Installs and enables a plugin.
	"""
	if not Plugin.install(args.plugin):
		return 1

	plugin = Plugin(args.plugin)
	plugin.enabled = True

###############################################################################
def enable_plugin(args):
	""" Enables a plugin.
	"""
	Plugin(args.plugin).enabled = True

###############################################################################
def disable_plugin(args):
	""" Disables a plugin.
	"""
	Plugin(args.plugin).enabled = False

###############################################################################
def list_plugin(args):						# pylint: disable=unused-argument
	""" Lists all enabled plugins.
	"""
	for plugin in Plugin.get_enabled_plugins():
		print(plugin)

###############################################################################
def build_parser():
	""" Constructs an argument parser and returns the parsed arguments.
	"""
	parser = argparse.ArgumentParser(
		description='Descriptive deep learning')
	parser.add_argument('--no-color', action='store_true',
		help='Disable colorful logging.')
	parser.add_argument('-v', '--verbose', default=0, action='count',
		help='Increase verbosity. Can be specified up to three times for '
			'trace-level output.')
	parser.add_argument('--monitor', action='store_true',
		help='Run Kur in monitor mode, which tries to recover from critical '
			'errors, like segmentation faults.')
	parser.add_argument('-V', '--version', action='store_true',
		help='Display version and exit.')
	parser.add_argument('--plugin', help='Plugin directory, or "none" to '
		'disable plugins.')

	subparsers = parser.add_subparsers(dest='cmd', help='Sub-command help.')

	subparser = subparsers.add_parser('train', help='Trains a model.')
	subparser.add_argument('--step', action='store_true',
		help='Interactive debug; prompt user before submitting each batch.')
	subparser.add_argument('kurfile', help='The Kurfile to use.')
	subparser.set_defaults(func=train)

	subparser = subparsers.add_parser('test', help='Tests a model.')
	subparser.add_argument('--step', action='store_true',
		help='Interactive debug; prompt user before submitting each batch.')
	subparser.add_argument('kurfile', help='The Kurfile to use.')
	subparser.set_defaults(func=test)

	subparser = subparsers.add_parser('evaluate', help='Evaluates a model.')
	subparser.add_argument('--step', action='store_true',
		help='Interactive debug; prompt user before submitting each batch.')
	subparser.add_argument('kurfile', help='The Kurfile to use.')
	subparser.set_defaults(func=evaluate)

	subparser = subparsers.add_parser('build',
		help='Tries to build a model. This is useful for debugging a model.')
	subparser.add_argument('kurfile', help='The Kurfile to use.')
	subparser.add_argument('-c', '--compile',
		choices=['none', 'train', 'test', 'evaluate', 'auto'], default='auto',
		help='Try to compile the specified variation of the model. If '
			'--compile=none, then it only tries to assemble the model, not '
			'compile anything. --compile=none implies --bare')
	subparser.add_argument('-b', '--bare', action='store_true',
		help='Do not attempt to load the data providers. In order for your '
			'model to build correctly with this option, you will need to '
			'specify shapes for all of your inputs.')
	subparser.set_defaults(func=build)

	subparser = subparsers.add_parser('dump',
		help='Dumps the Kurfile out as a JSON blob. Useful for debugging.')
	subparser.add_argument('kurfile', help='The Kurfile to use.')
	subparser.add_argument('-p', '--pre-parse', action='store_true',
		help='Dump the Kurfile before parsing it.')
	subparser.set_defaults(func=dump)

	subparser = subparsers.add_parser('data',
		help='Does not actually compile anything, but only prints out a '
			'single batch of data. This is useful for debugging data sources.')
	subparser.add_argument('kurfile', help='The Kurfile to use.')
	subparser.add_argument('-t', '--target',
		choices=['train', 'validate', 'test', 'evaluate', 'auto'],
		default='auto', help='Try to produce data corresponding to a specific '
			'variation of the model.')
	subparser.add_argument('--assemble', action='store_true', help='Also '
		'begin assembling the model to pull in compile-time, auxiliary data '
		'sources.')
	subparser.add_argument('-n', '--number', type=int,
		help='Number of samples to print (default: the entire batch).')
	subparser.set_defaults(func=prepare_data)

	###########################################################################
	# Plugins

	subparser = subparsers.add_parser('plugin', help='Configures Kur plugins.')
	subsubparsers = subparser.add_subparsers(dest='plugin_cmd',
		help='Plugin command')

	subsubparser = subsubparsers.add_parser('install', help='Install a plugin.')
	subsubparser.add_argument('plugin', help='The plugin to install.')
	subsubparser.set_defaults(func=install_plugin)

	subsubparser = subsubparsers.add_parser('enable', help='Enable a plugin.')
	subsubparser.add_argument('plugin', help='The plugin to enable.')
	subsubparser.set_defaults(func=enable_plugin)

	subsubparser = subsubparsers.add_parser('disable',
		help='Disable a plugin.')
	subsubparser.add_argument('plugin', help='The plugin to disable.')
	subsubparser.set_defaults(func=disable_plugin)

	subsubparser = subsubparsers.add_parser('list',
		help='List enabled plugin.')
	subsubparser.set_defaults(func=list_plugin)

	return parser, subparsers

###############################################################################
def parse_args(parser):
	""" Parses command-line arguments.
	"""
	return parser.parse_args()

###############################################################################
def main():
	""" Entry point for the Kur command-line script.
	"""
	gotcha = False
	plugin_dir = None
	for arg in sys.argv[1:]:
		if gotcha:
			plugin_dir = arg
			break
		elif arg == '--plugin':
			gotcha = True
	plugin_dir = plugin_dir or os.environ.get('KUR_PLUGIN')
	load_plugins(plugin_dir)

	parser, _ = build_parser()
	args = parse_args(parser)

	loglevel = {
		0 : logging.WARNING,
		1 : logging.INFO,
		2 : logging.DEBUG,
		3 : logging.TRACE
	}
	config = logging.basicConfig if args.no_color else logcolor.basicConfig
	config(
		level=loglevel.get(args.verbose, logging.TRACE),
		format='{color}[%(levelname)s %(asctime)s %(name)s:%(lineno)s]{reset} '
			'%(message)s'.format(
				color='' if args.no_color else '$COLOR',
				reset='' if args.no_color else '$RESET'
			)
	)
	logging.captureWarnings(True)

	do_monitor(args)

	if args.version:
		args.func = version
	elif not hasattr(args, 'func'):
		print('Nothing to do!', file=sys.stderr)
		print('For usage information, try: kur --help', file=sys.stderr)
		print('Or visit our homepage: {}'.format(__homepage__))
		sys.exit(1)

	engine = JinjaEngine()
	setattr(args, 'engine', engine)

	sys.exit(args.func(args) or 0)

###############################################################################
if __name__ == '__main__':
	main()

### EOF.EOF.EOF.EOF.EOF.EOF.EOF.EOF.EOF.EOF.EOF.EOF.EOF.EOF.EOF.EOF.EOF.EOF.EOF
