#!/usr/bin/env python3
import argparse
import configparser
import difflib
import io
import os
import re
import requests, requests.exceptions
import shutil
import subprocess
import sys
import time
import webbrowser
import zipfile

ANSI_RED = "\033[31m"
ANSI_GREEN = "\033[32m"
ANSI_YELLOW = "\033[33m"
ANSI_BLUE = "\033[34m"
ANSI_RESET = "\033[m"

def relativePath(path):
	abspath = os.path.abspath(__file__)
	dirname = os.path.dirname(abspath)
	return os.path.join(dirname, path)

def declare(messages):
	if type(messages) is str:
		print(f"{ANSI_YELLOW}{messages}{ANSI_RESET}")
	else:
		for i in messages:
			print(f"{ANSI_YELLOW}{i}{ANSI_RESET}")

def say(messages):
	if type(messages) is str:
		messages = messages.splitlines()
	for i in messages:
		for j in i.splitlines():
			print(f"\t{j}")

def pester(messages):
	if type(messages) is str:
		messages = messages.splitlines()
	for i in messages:
		for j in i.splitlines():
			print(f"\t{j}")

def reset():
	print("\033[F\033[K", end='')

def lastLang():
	langfile = relativePath(".lang")
	if os.path.exists(langfile):
		with open(langfile) as f:
			return f.readline().strip()
	return "py"

DEFAULT_CONFIG = '/usr/local/etc/kattisrc'
HEADERS = {'User-Agent': 'kattis-cli-submit'}
EXTENSIONS = {'py': 'Python 3', 'rb': 'Ruby', 'hs': 'Haskell'}
LANG = lastLang()
INTERPRETER = {'py': sys.executable, 'rb': 'ruby', 'hs': 'runhaskell'}

class ConfigError(Exception):
	pass

def getUrl(cfg, option, default):
	if cfg.has_option('kattis', option):
		return cfg.get('kattis', option)
	else:
		return 'https://%s/%s' % (cfg.get('kattis', 'hostname'), default)

def getConfig():
	cfg = configparser.ConfigParser()
	if os.path.exists(DEFAULT_CONFIG):
		cfg.read(DEFAULT_CONFIG)
	if not cfg.read([os.path.join(os.getenv('HOME'), '.kattisrc'),
					os.path.join(os.path.dirname(sys.argv[0]), '.kattisrc')]):
		raise ConfigError('''\
I failed to read in a config file from your home directory or from the
same directory as this script. To download a .kattisrc file please visit
https://<kattis>/download/kattisrc
The file should look something like this:
[user]
username: yourusername
token: *********
[kattis]
hostname: <kattis>
loginurl: https://<kattis>/login
submissionurl: https://<kattis>/submit
submissionsurl: https://<kattis>/submissions''')
	return cfg

def login(loginUrl, username, password=None, token=None):
	loginArgs = {'user': username, 'script': 'true'}
	if password:
		loginArgs['password'] = password
	if token:
		loginArgs['token'] = token
	return requests.post(loginUrl, data=loginArgs, headers=HEADERS)

def loginFromConfig(cfg):
	username = cfg.get('user', 'username')
	fields = ['password', 'token']
	data = [None, None]
	for i in range(2):
		try:
			data[i] = cfg.get('user', fields[i])
		except configparser.NoOptionError:
			pass
	if all(v is None for v in data):
		raise ConfigError('''\
Your .kattisrc file appears corrupted. It must provide a token (or a
KATTIS password).
Please download a new .kattisrc file''')
	loginUrl = getUrl(cfg, 'loginurl', 'login')
	return login(loginUrl, username, data[0], data[1])

def requestSubmit(submitUrl, cookies, problem, filename):
	global LANG
	data = {'submit': 'true',
			'submit_ctr': 2,
			'language': EXTENSIONS[LANG],
			'mainclass': '',
			'problem': problem,
			'tag': '',
			'script': 'true'}
	subFiles = []
	with open(filename) as subFile:
		subFiles.append(('sub_file[]',
			(os.path.basename(filename), subFile.read(), 'application/octet-stream')))
	return requests.post(submitUrl, data=data, files=subFiles, cookies=cookies, headers=HEADERS)

def getCookies():
	try:
		cfg = getConfig()
	except ConfigError as exc:
		pester(exc)
		sys.exit(1)

	try:
		loginReply = loginFromConfig(cfg)
	except ConfigError as exc:
		pester(exc)
		sys.exit(1)
	except requests.exceptions.RequestException as err:
		pester(f'Login connection failed: {err}')
		sys.exit(1)
	if not loginReply.status_code == 200:
		pester('Login failed.')
		pester(f'Status code: {loginReply.status_code}')
		sys.exit(1)
	return loginReply.cookies

def submit(problem, file):
	try:
		cfg = getConfig()
	except ConfigError as exc:
		pester(exc)
		sys.exit(1)

	submitUrl = getUrl(cfg, 'submissionurl', 'submit')
	cookies = getCookies()
	try:
		result = requestSubmit(submitUrl, cookies, problem, file)
	except requests.exceptions.RequestException as err:
		pester('Submit connection failed:', err)
		sys.exit(1)
	if result.status_code != 200:
		pester('Submission failed.')
		pester(f'Status code: {loginReply.status_code}')
		sys.exit(1)
	plainResult = result.content.decode('utf-8').replace('<br />', '\n')
	say(plainResult)
	pollSubmitStatus(plainResult, cfg)

def pollSubmitStatus(submitResponse, cfg):
	pattern = r'data-type="status"><span class="([^"]+).*</i>([^<]+)'
	correct = r'span class="accepted" title="Test case'
	incorrect = r'span class="rejected" title="Test case'
	unchecked = r'Test case \d+/\d+: not checked'

	submissionUrl = getUrl(cfg, 'submissionsurl', 'submissions')
	match = re.search(r'Submission ID: (\d+)', submitResponse)
	if match:
		id = match.group(1)
		resultUrl = f"{submissionUrl}/{id}"
		cookies = getCookies()
		say("Kattis Tests: []")
		while True:
			try:
				result = requests.get(resultUrl, cookies=cookies)
			except requests.exceptions.RequestException as err:
				pester('Result page request failed:', err)
				sys.exit(1)
			if result.status_code != 200:
				pester('Result page request failed.')
				pester(f'Status code: {result.status_code}')
				sys.exit(1)
			plainResult = result.content.decode('utf-8')
			match2 = re.search(pattern, plainResult)
			if match2:
				correctAmount = len(re.findall(correct, plainResult))
				incorrectAmount = len(re.findall(incorrect, plainResult))
				uncheckedAmount = len(re.findall(unchecked, plainResult))
				progress = f"{ANSI_GREEN}={ANSI_RESET}" * correctAmount
				progress += f"{ANSI_RED}x{ANSI_RESET}" * incorrectAmount
				progress += "." * uncheckedAmount
				reset()
				say(f"Kattis Tests: [{progress}]")

				status1 = match2.group(1)
				status2 = match2.group(2)
				if status1 != 'other':
					say(f"Result: {status2}")
					break
			else:
				break
			time.sleep(0.33)

def testProblem(problem):
	global LANG
	folder = relativePath("problems")
	problemDir = os.path.join(folder, problem)
	script = os.path.join(problemDir, f"solution.{LANG}")
	if not os.path.exists(script):
		pester(f"{problem}/solution.{LANG} doesn't exist")
		return
	testDir = os.path.join(problemDir, "tests")
	if not os.path.exists(testDir):
		pester(f"{problem}/tests/ doesn't exist")

	testFiles = [t for t in os.listdir(testDir)
					if os.path.isfile(os.path.join(testDir, t))]
	inputs = {}
	outputs = {}
	for filename in testFiles:
		parts = filename.split(".")
		head = ".".join(parts[:-1])
		if parts[-1] == "in":
			inputs[head] = os.path.join(testDir, filename)
		elif parts[-1] == "ans":
			outputs[head] = os.path.join(testDir, filename)
	errors = []
	numTests = len([x for x in inputs.keys() if x in outputs])
	testsDone = 0
	progress = ""

	say("Sample Tests: [" + ('.'*numTests) + "]")
	for i, infile in inputs.items():
		if i not in outputs:
			continue
		outfile = outputs[i]
		process = subprocess.run([INTERPRETER[LANG], script],
						capture_output=True,
						stdin=open(infile))
		out, err = process.stdout.decode("utf-8"), process.stderr.decode("utf-8")
		if process.returncode != 0:
			errors += [' ', f"FAIL: {infile}", '-'*(6+len(infile))]
			errors.append(f"Nonzero exit code: {process.returncode}")
			errors.append(f"stderr: {err}")
			progress += f"{ANSI_RED}x{ANSI_RESET}"
		else:
			expected = open(outfile).read().splitlines(keepends=True)
			actual = out.splitlines(keepends=True)
			diff = list(difflib.Differ().compare(expected, actual))
			if len(diff) == len(expected) and all([line.startswith("  ") for line in diff]):
				progress += f"{ANSI_GREEN}={ANSI_RESET}"
			else:
				progress += f"{ANSI_RED}x{ANSI_RESET}"
				errors += [' ', f"FAIL: {infile}", '-'*(5+len(infile))]
				errors.append("Expected - but received +")
				errors += diff
		testsDone += 1
		reset()
		say(f"Sample Tests: [" + progress + ('.'*(numTests-testsDone)) + "]")
	pester(errors)

def boilerplate(problem):
	global LANG
	folder = relativePath("problems")
	problemDir = os.path.join(folder, problem)
	filename = os.path.join(problemDir, f"solution.{LANG}")
	templateFolder = relativePath(f"templates/{LANG}")
	if not os.path.exists(templateFolder):
		pester("No template folder found at: "+templateFolder)
		return

	needed_extension = f".{LANG}"
	templates = [f[:-len(needed_extension)] for f in os.listdir(templateFolder)
				if os.path.isfile(os.path.join(templateFolder, f))
					and f.endswith(needed_extension)]

	options = ', '.join([k for k in templates])
	print(f"{ANSI_BLUE}>> Input one of ({options}):{ANSI_RESET} ", end='')
	sys.stdout.flush()
	choice = input().strip()
	if choice in templates:
		declare(f"boilerplate {choice}")
		templateFile = os.path.join(templateFolder, choice + needed_extension)
		originalText = ""
		templateText = ""
		with open(templateFile, 'r') as contents:
			templateText = contents.read()
		try:
			with open(filename, 'r') as contents:
				originalText = contents.read()
		except FileNotFoundError:
			pass
		with open(filename, 'w') as contents:
			contents.write(templateText)
			contents.write(originalText)
	else:
		print(f"{ANSI_RED}invalid input:{ANSI_RESET} {choice}")
		sys.stdout.flush()

def setupProblemEnvironment(problem):
	global LANG
	folder = relativePath("problems")
	problemDir = os.path.join(folder, problem)
	testDir = os.path.join(problemDir, "tests")
	filename = os.path.join(problemDir, f"solution.{LANG}")
	if not os.path.exists(testDir):
		os.makedirs(testDir, exist_ok=True)
		open(filename, 'w').close()
		zipUrl = f"https://open.kattis.com/problems/{problem}/file/statement/samples.zip"
		zipResponse = requests.get(zipUrl)
		try:
			testData = zipfile.ZipFile(io.BytesIO(zipResponse.content))
		except zipfile.BadZipFile:
			pester("Problem not found.")
			shutil.rmtree(problemDir)
			sys.exit(1)
		testData.extractall(testDir)
	return filename

def openEditor(filename):
	editor = os.environ.get("EDITOR", "vim")
	subprocess.call([editor, filename])

def main():
	global LANG
	parser = argparse.ArgumentParser(description='Solve, test, submit Kattis problems.')
	parser.add_argument('problem', nargs=1, help='Kattis problem id')
	args = parser.parse_args()

	problem = args.problem[0]
	filename = setupProblemEnvironment(problem)

	options = {'e': 'edit', 'b': 'boilerplate', 'l': 'language', 't': 'test', 's': 'submit', 'q': 'quit'}
	optiontext = ", ".join([f"{k} => {v}" for k, v in options.items()])

	while True:
		print(f"{ANSI_BLUE}>> Input one of ({optiontext}):{ANSI_RESET} ", end='')
		sys.stdout.flush()
		try:
			key = input().strip().lower()
		except EOFError:
			sys.exit()
		reset()

		if key not in options:
			print(f"{ANSI_RED}invalid input:{ANSI_RESET} {key}")
			continue

		action = options[key]
		if action == "quit":
			declare("quit")
			break
		elif action == "boilerplate":
			boilerplate(problem)
		elif action == "edit":
			declare(f"edit {filename}")
			openEditor(filename)
		elif action == "language":
			optionlang = ', '.join([k for k in EXTENSIONS])
			print(f"{ANSI_BLUE}>> Input one of ({optionlang}):{ANSI_RESET} ", end='')
			sys.stdout.flush()
			choice = input().strip().lower()
			if choice in EXTENSIONS:
				LANG = choice
				filename = setupProblemEnvironment(problem)
				declare(f"lang {choice}")
				langfile = relativePath(".lang")
				with open(langfile, "w") as f:
					f.write(choice)
				f.close()
			else:
				print(f"{ANSI_RED}invalid input:{ANSI_RESET} {choice}")
				sys.stdout.flush()
		elif action == "submit":
			declare("submit")
			submit(problem, filename)
		elif action == "test":
			declare(f"test {filename}")
			testProblem(problem)
			pass

if __name__ == '__main__':
	main()