import sys
import tempfile
import os.path
import os
import shutil

print(sys.argv[-1])

if sys.argv[-1] == 'data.py' or 'data.yaml' not in sys.argv[-1]:
	print('Please provide a path to a data.yaml file')
	exit(1)

filePath = sys.argv[-1]
data = []
basePath = os.path.dirname(filePath)
trainPath = os.path.join(basePath, 'train', 'images')
validPath = os.path.join(basePath, 'valid', 'images')
testPath = os.path.join(basePath, 'test', 'images')

if not os.path.exists(basePath):
	print(f'Dataset directory {basePath} doesn\'t exist')
	exit(2)

if not os.path.exists(trainPath):
	print(f'Dataset directory {trainPath} doesn\'t exist')
	exit(2)

if not os.path.exists(validPath):
	print(f'Dataset directory {validPath} doesn\'t exist')
	exit(2)

if not os.path.exists(testPath):
	print(f'Dataset directory {testPath} doesn\'t exist')
	exit(2)

def isInList (match, l) :
	for i in l:
		if match in i:
			return True
	return False

with open(filePath, 'r') as file:
	lines = file.readlines()
	for line in lines:
		if 'train: ' in line:
			data.append(f'{os.linesep}')
			data.append(f'train: {trainPath}{os.linesep}')
		elif 'val: ' in line:
			data.append(f'{os.linesep}')
			data.append(f'val: {validPath}{os.linesep}')
		elif 'test: ' in line:
			data.append(f'{os.linesep}')
			data.append(f'test: {testPath}{os.linesep}')
		else:
			data.append(line)
	if not isInList('train: ', data):
		data.append(f'{os.linesep}')
		data.append(f'train: {trainPath}{os.linesep}')
	elif not isInList('val: ', data):
		data.append(f'{os.linesep}')
		data.append(f'val: {validPath}{os.linesep}')
	elif not isInList('test: ', data):
		data.append(f'{os.linesep}')
		data.append(f'test: {testPath}{os.linesep}')
	file.close()

with tempfile.NamedTemporaryFile() as tmp:
	with open(tmp.name, 'w') as t:
		t.writelines(data)
	shutil.copy(tmp.name, filePath)