import os,sys
import subprocess
import time
from termcolor import colored
from rich.console import Console
from rich.table import Table

def add_timestamp(log_file, attrib, task_name):
	with open(log_file, 'a') as f: f.write(f"{task_name}\t{attrib}\t{int(time.time())}\n")

def parse_log(log_path, fill_attribs = True):
	log = {}
	if os.path.exists(log_path):
		for entry in os.listdir(log_path):
			if entry == '.' or entry == '..': continue
			with open(os.path.join(log_path,entry)) as f:
				for line in f:
					name, status, time_int = line.rstrip().split("\t")
					time = int(time_int)
					query = log.get(name)
					if query == None:
						log[name] = {status: [time]}
					else:
						query_status = query.get(status)
						if query_status == None: 
							query[status] = [time]
						else:
							query[status].append(time)
	if fill_attribs:
		for task, attribs in log.items():
			#puts "#{attribs.inspect}"
			set_length = len(attribs['set'])
			fill_attrib(attribs, 'start', set_length)
			fill_attrib(attribs, 'end', set_length)
			fill_attrib(attribs, 'error', set_length)
	return log

def fill_attrib(attribs, mode, set_length):
	query = attribs.get(mode)
	if query == None:
		attribs[mode] = [ 0 for x in range(set_length)]
	elif len(query) < set_length:
		for x in range(set_length -len(query)): query.append(0)

def report_log(log, initial_flow_attribs, mode, workflow_status, no_size, raw):
	set_task_state(log, workflow_status)
	set_time(log)
	if mode == None or mode.upper() == 'ALL':
		tasks = log
	else:
		tasks = {}
		for name, attribs in log.items(): 
			if attribs['state'] == mode.upper(): tasks[name] = attribs
	rows = []
	for  task_name, attribs in tasks.items():
		job_path = initial_flow_attribs[task_name][0]
		size = None
		if not no_size: size = subprocess.check_output(f"du -sh {job_path}", shell=True).decode("utf-8").split()[0] 
		rows.append([attribs['state_msg'], os.path.basename(job_path), attribs['time'], size, task_name])
	
	header = ['Status', 'Folder', 'Time', 'Size', 'Job Name']
	if raw:
		rows.insert(0, header)
		formated_report = rows
	else:
		table = Table(title='Logging')
		for col in header: table.add_column(col)
		for row in rows: table.add_row(*row)
		console = Console(record=True)
		console.print(table)
		formated_report = console
	return formated_report

def set_task_state(log, workflow_status, position = -1):
	for task, attribs in log.items():
		start_position = len(attribs['start']) - position
		start = attribs['start'][position]
		stop_position = len(attribs['end']) - position
		stop = attribs['end'][position]
		if attribs['error'][position] > 0:
			status = 'ERROR'
			status_msg = colored('ERROR', 'red')
		elif workflow_status: # Workflow has finished
			if start == 0 and stop == 0:
				status = 'NOT'
				status_msg = colored('NOT', 'blue')
			elif start > 0 and stop > 0:
				status = 'SUCC'
				status_msg = colored('SUCC', 'green')
			elif start > 0 and stop ==0:
				status = 'ABORT'
				status_msg = colored('ABORT', 'red')
		else: # Workflow is still running
			if start == 0 and stop == 0:
				status = 'PEND'
				status_msg = colored('PEND', 'blue')
			elif start > 0 and stop > 0:
				status = 'SUCC'
				status_msg = colored('SUCC', 'green')
			elif start > 0 and stop == 0:
				status = 'RUN'
				status_msg = colored('RUN', 'magenta')
		attribs['state'] = status
		attribs['state_msg'] = status_msg

def set_time(log):
	for task, attribs in log.items():
		start = attribs['start'][-1]
		stop = attribs['end'][-1]
		error = attribs['error'][-1]
		status = attribs['state']
		timestamp = 0
		if status == 'SUCC':
			timestamp = stop - start
		elif status == 'RUN':
			timestamp = int(time.time()) - start
		elif status == 'ERROR':
			timestamp = error - start
		attribs['seconds'] = timestamp
		magnitude = 's'
		if timestamp >= 60:
			magnitude = 'm'
			timestamp = timestamp/60.0 # To minutes
		if timestamp >= 60 and magnitude == 'm':
			magnitude = 'h'
			timestamp = timestamp/60.0 # To hours
		if timestamp >= 24 and magnitude == 'h':
			magnitude = 'd'
			timestamp = timestamp/24.0 # To days
		if timestamp == 0:
			time_string = '-'
		else:
			time_string = f"{timestamp} {magnitude}"
		attribs['time'] = time_string

def write_log(log, log_path, job_relations_with_folders, write_task_logs= True):
	if not os.path.exists(log_path): os.mkdir(log_path)
	if write_task_logs:
		for name, folder_deps in job_relations_with_folders.items():
			if log.get(name) != None: #Control check when the wk_log folder has been deleted
				folder, deps = folder_deps
				with open(('/').join([log_path, os.path.basename(folder)]), 'w') as f:
					for mode, times in log[name].items():
						for time in times: 
							f.write(f"{name}\t{mode}\t{time}\n")
