#!/usr/bin/python3
# SPDX-License-Identifier: MIT

import argparse

import colorama
import sys
import xbstrap.base
import yaml

main_parser = argparse.ArgumentParser()
main_parser.add_argument('-v', dest='verbose', action='store_true',
		help="verbose")
main_subparsers = main_parser.add_subparsers(dest='command')

class Pipeline:
	def __init__(self, cfg, pipe_yml):
		self.cfg = cfg
		self.jobs = dict()

		# Determine the set of jobs.
		mentioned_tools = set()
		mentioned_pkgs = set()
		for job_yml in pipe_yml['jobs']:
			tools = []
			pkgs = []
			for name in job_yml.get('tools', []):
				tool = cfg.get_tool_pkg(name)
				tools.append(tool)
				mentioned_tools.add(tool)
			for name in job_yml.get('packages', []):
				pkg = cfg.get_target_pkg(name)
				pkgs.append(pkg)
				mentioned_pkgs.add(pkg)

			name = job_yml['name']
			assert name not in self.jobs
			self.jobs[name] = Job(name, tools, pkgs)

		for tool in cfg.all_tools():
			if tool in mentioned_tools:
				continue
			if not tool.is_default:
				continue
			name = 'tool:' + tool.name
			assert name not in self.jobs
			self.jobs[name] = Job(name, [tool], [])
		for pkg in cfg.all_pkgs():
			if pkg in mentioned_pkgs:
				continue
			if not pkg.is_default:
				continue
			name ='package:' + pkg.name
			assert name not in self.jobs
			self.jobs[name] = Job(name, [], [pkg])

	def all_jobs(self):
		return self.jobs.values()

	def get_job(self, name):
		return self.jobs[name]

class Job:
	def __init__(self, name, tools, pkgs):
		self.name = name
		self.tools = set(tools)
		self.pkgs = set(pkgs)

def pipeline_for_dir(cfg):
	with open('pipeline.yml', 'r') as f:
		pipe_yml = yaml.load(f, yaml.SafeLoader)
	return Pipeline(cfg, pipe_yml)

class PipelineItem:
	def __init__(self, job):
		self.job = job
		self.edge_set = set()
		self.edge_list = []
		self.plan_state = xbstrap.base.PlanState.NULL
		self.resolved_n = 0

def do_compute_graph(args):
	cfg = xbstrap.base.config_for_dir()
	pipe = pipeline_for_dir(cfg)

	items = dict()
	for job in pipe.all_jobs():
		item = PipelineItem(job)
		items[item.job.name] = item

	tool_mapping = dict()
	pkg_mapping = dict()
	for item in items.values():
		for tool in item.job.tools:
			tool_mapping[tool] = item.job.name
		for pkg in item.job.pkgs:
			pkg_mapping[pkg] = item.job.name

	for item in items.values():
		plan = xbstrap.base.Plan(cfg)
		plan.build_scope = set().union(item.job.tools, item.job.pkgs)
		for tool in item.job.tools:
			plan.wanted.update([(xbstrap.base.Action.INSTALL_TOOL_STAGE, stage)
					for stage in tool.all_stages()])
		for pkg in item.job.pkgs:
			if cfg.use_xbps:
				plan.wanted.update([(xbstrap.base.Action.PACK_PKG, pkg)])
			else:
				plan.wanted.update([(xbstrap.base.Action.BUILD_PKG, pkg)])
		plan.compute_plan(no_ordering=True)

		for (action, subject) in plan.materialized_steps():
			if action == xbstrap.base.Action.WANT_TOOL:
				if subject in item.job.tools:
					continue
				item.edge_set.add(tool_mapping[subject])
			if action == xbstrap.base.Action.WANT_PKG:
				if subject in item.job.pkgs:
					continue
				item.edge_set.add(pkg_mapping[subject])

	for item in items.values():
		item.edge_list = list(item.edge_set)

	order = []

	# TODO: this is copied from the planning code. Unify these code paths!
	# The following code does a topologic sort of the desired items.
	stack = []

	def visit(item):
		if item.plan_state == xbstrap.base.PlanState.NULL:
			item.plan_state = xbstrap.base.PlanState.EXPANDING
			stack.append(item)
		elif item.plan_state == xbstrap.base.PlanState.EXPANDING:
			reverse_chain = [item]
			for circ_item in reversed(stack):
				reverse_chain.append(circ_item)
				if circ_item == item:
					break
			chain = reversed(reverse_chain)
			raise RuntimeError("Job has circular dependencies {}".format(
					[chain_item.job.name for chain_item in chain]))
		else:
			# Packages that are already ordered do not need to be considered again.
			assert item.plan_state == xbstrap.base.PlanState.ORDERED

	for root_item in items.values():
		visit(root_item)

		while stack:
			item = stack[-1]
			if item.resolved_n == len(item.edge_list):
				assert item.plan_state == xbstrap.base.PlanState.EXPANDING
				item.plan_state = xbstrap.base.PlanState.ORDERED
				stack.pop()
				order.append(item)
			else:
				edge_item = items[item.edge_list[item.resolved_n]]
				item.resolved_n += 1
				visit(edge_item)

	if args.gv:
		# For visualization purposes.
		print('digraph {');
		for item in order:
			for edge in item.edge_list:
				print('    "{}" -> "{}";'.format(edge, item.job.name))
		print('}')
	elif args.linear:
		for item in order:
			print('{}'.format(item.job.name))
	else:
		for item in order:
			print('{} {}'.format(item.job.name, ' '.join(item.edge_list)))

do_compute_graph.parser = main_subparsers.add_parser('compute-graph')
do_compute_graph.parser.add_argument('--linear', action='store_true')
do_compute_graph.parser.add_argument('--gv', action='store_true')

def do_run_job(args):
	cfg = xbstrap.base.config_for_dir()
	pipe = pipeline_for_dir(cfg)
	job = pipe.get_job(args.job)

	plan = xbstrap.base.Plan(cfg)
	if args.dry_run:
		plan.dry_run = True
	if args.check:
		plan.check = True
	if args.keep_going:
		plan.keep_going = True
	plan.build_scope = set().union(job.tools, job.pkgs)
	for tool in job.tools:
		plan.wanted.update([(xbstrap.base.Action.INSTALL_TOOL_STAGE, stage)
				for stage in tool.all_stages()])
	for pkg in job.pkgs:
		if cfg.use_xbps:
			plan.wanted.update([(xbstrap.base.Action.PACK_PKG, pkg)])
		else:
			plan.wanted.update([(xbstrap.base.Action.BUILD_PKG, pkg)])
	plan.run_plan()

do_run_job.parser = main_subparsers.add_parser('run-job')
do_run_job.parser.add_argument('job', type=str)
do_run_job.parser.add_argument('-n', '--dry-run', action='store_true',
		help="compute a plan but do not execute it")
do_run_job.parser.add_argument('-c', '--check', action='store_true',
		help="skip packages that are already built/installed/etc.")
do_run_job.parser.add_argument('--keep-going', action='store_true',
		help="continue running even if some build steps fail")

args = main_parser.parse_args()

colorama.init()

if args.verbose:
	xbstrap.base.verbosity = True

try:
	if args.command == 'compute-graph':
		do_compute_graph(args)
	elif args.command == 'run-job':
		do_run_job(args)
	else:
		assert not "Unexpected command"
except (xbstrap.base.ExecutionFailureException, xbstrap.base.PlanFailureException) as e:
	print('{}xbstrap{}: {}{}{}'.format(colorama.Style.BRIGHT, colorama.Style.RESET_ALL,
			colorama.Fore.RED, e, colorama.Style.RESET_ALL))
	sys.exit(1)
except KeyboardInterrupt as e:
	sys.exit(1)
