#!/usr/bin/env python3
import argparse
import os
import sys
import mssh
from shutil import copyfile
from mssh.config import MsshConfig,get_cnf_file_path
import getpass

cnf = MsshConfig().load()


def completions(args):
	try:
		if args.env is None:
			print('\n'.join(cnf.environments.keys()))
			return
		else:
			if args.env not in cnf.environments:
				print('\n'.join([e for e in cnf.environments.keys() if e.startswith(args.env)]))
				return

		if args.product is None:
			print('\n'.join(cnf.environments[args.env].keys()))
			return
		else:
			if args.product not in cnf.environments[args.env]:
				print('\n'.join([e for e in cnf.environments[args.env].keys() if e.startswith(args.product)]))
				return

		if args.server is None:
			print('\n'.join(cnf.environments[args.env][args.product]['servers'].keys()))
			return
		else:
			if args.server not in cnf.environments[args.env][args.product]['servers']:
				print('\n'.join([e for e in cnf.environments[args.env][args.product]['servers'].keys() if e.startswith(args.server)]))
				return
				
		if args.user is None:
			print('\n'.join(cnf.environments[args.env][args.product]['servers'][args.server]['users']))
			return
		else:
			if args.user not in cnf.environments[args.env][args.product]['servers'][args.server]['users']:
				print('\n'.join([e for e in cnf.environments[args.env][args.product]['servers'][args.server]['users'] if e.startswith(args.user)]))
				return

		if args.command is None or len(args.command)==0:
			print('\n'.join(cnf.environments[args.env][args.product]['servers'][args.server]['commands'].keys()))
			return
		else:
			if len(args.command)==1 and args.command[0] not in cnf.environments[args.env][args.product]['servers'][args.server]['commands']:
				print('\n'.join([e for e in cnf.environments[args.env][args.product]['servers'][args.server]['commands'].keys() if e.startswith(args.command[0])]))
				return

	except KeyError:
		return

def edit_config():
	os.system("xdg-open " + get_cnf_file_path())

def setup():
	if getpass.getuser() != 'root':
		print("You have to run this as root", file=sys.stderr)
		return
		
	bash_completion_system_path = '/usr/share/bash-completion/completions'
	if os.path.exists(bash_completion_system_path):
		bash_installation = input ("Would you like to install completions for bash? [Y/n]: ")
		if bash_installation=="" or bash_installation=='y' or bash_installation=='Y':
			bash_completion_path = os.path.join(mssh.__path__[0], 'completions','bash','mssh')
			copyfile(bash_completion_path,bash_completion_system_path+'/mssh')
	
	zsh_completion_system_path='/usr/share/zsh/vendor-completions'
	if os.path.exists(zsh_completion_system_path):
		zsh_installation = input ("Would you like to install completions for zsh? [Y/n]: ")
		if zsh_installation=="" or zsh_installation=='y' or zsh_installation=='Y':
			zsh_completion_path = os.path.join(mssh.__path__[0], 'completions','zsh','_mssh')
			copyfile(zsh_completion_path,zsh_completion_system_path+'/_mssh')

def main():
	parser = argparse.ArgumentParser(description='Ssh Manager')
	parser.add_argument('--completions', action='store_true', help=argparse.SUPPRESS)
	parser.add_argument('--setup', action='store_true', help='Setup auto completions')
	parser.add_argument('--edit', action='store_true', help='Edit configuration')
	parser.add_argument('env', help='Environment', nargs='?', type=str, default=None)
	parser.add_argument('product', help='Product name', type=str, nargs='?', default=None)
	parser.add_argument('server', help='Server name or alias', nargs='?', type=str, default=None)
	parser.add_argument('user', help='User', type=str, nargs='?', default=None)
	parser.add_argument('command', help='Command', type=str, nargs=argparse.REMAINDER, default=None)
	args = parser.parse_args(sys.argv[1:])

	if args.completions:
		completions(args)
		return 
	if args.setup:
		setup()
		return
	if args.edit:
		edit_config()
		return

	if not (args.env in cnf.environments):
		print("Environment %s not found" % str(args.env), file=sys.stderr)
		sys.exit(1)
	if args.product is None:
		print("Products in [%s] environment" % args.env)
		for prod in cnf.environments[args.env].keys():
			print(prod)
	else:
		if not (args.product in cnf.environments[args.env]):
			print("Product %s not found" % str(args.product), file=sys.stderr)
			sys.exit(1)
		conn_server = args.server
		if args.server is None:
			if len(cnf.environments[args.env][args.product]['servers']) == 1:
				conn_server = next(iter(cnf.environments[args.env][args.product]['servers'].keys()))
			else:
				print("Servers in [%s] environment on [%s] product" % (args.env, args.product))
				for srv_name, srv in cnf.environments[args.env][args.product]['servers'].items():
					print("Name:%s Host:%s Users:%s" % (srv_name, srv.host, srv.users))
					sys.exit(1)

		if not (conn_server in cnf.environments[args.env][args.product]['servers']):
			print("Server %s not found" % str(args.server), file=sys.stderr)
			sys.exit(1)
		try:
			srv = cnf.environments[args.env][args.product]['servers'][conn_server]
			conn_user = args.user
			if args.user is None and len(srv.users) == 1:
				conn_user = srv.users[0]

			if conn_user is not None:
				cmd = "ssh %s@%s" % (conn_user, srv.host)
				
				if args.command is not None and len(args.command)>0:
					if len(args.command)>1:
						cmd += ' ' + ' '.join(args.command)
					else:
						# search command alias
						curr_server = cnf.environments[args.env][args.product]['servers'][conn_server]
						if 'commands' in curr_server and args.command[0] in curr_server['commands']:
							cmd += ' "'+cnf.environments[args.env][args.product]['servers'][conn_server]['commands'][args.command[0]]+'"'
						else:
							cmd += ' "' + args.command[0]+'"'

				print(cmd)	
				os.system(cmd)
			else:
				print("User parameter is mandatory", file=sys.stderr)
				sys.exit(1)

		except KeyError as e:
			print("Environment %s not found" % str(e), file=sys.stderr)


if __name__ == '__main__':
	main()
