#!/usr/bin/env python4

import re
import argparse
import io
import os

from sys import argv, exit


parser = argparse.ArgumentParser()

parser.add_argument("-m", "--manual", dest = "manual", default = None, help="Manual Mode {node #}")
parser.add_argument("-s", "--state_info", dest="state_info", help="State Info File")
parser.add_argument("-i", "--input", dest = "input_d", help="Input Dot File")
parser.add_argument("-o", "--output", dest = "output_d", help="Output Dot File.")


args=parser.parse_args()
#if len(argv) != 4:
 #   print("Usage: %s [node ID to path walk from] [input DOT file] [output DOT file]" %argv[0])
  #  exit(1)
if (args.manual == "0"):
    print("This is the root!")
    exit(1)

crit_nodes = []
if(args.manual !=None):
    crit_nodes.append(args.manual)
else:
    with open(args.state_info, 'r') as myfile:
        text = myfile.read()
    splits = text.split("ID:")
    states = len(splits) -1

    for i in range(states):
        if "is_critical=true" in splits[i]:
            crit_nodes.append(str(i-1))
            print(str(i-1))

for crits in crit_nodes:

    with open(args.input_d, 'r') as in_file:
        text=in_file.read()
    in_file.close()

    # Splits input file line by line, except for first and last line
    line_splits = text.split(";")
    total_lines = len(line_splits)

    # The first part of the file - contains the declaration of each state and any coloring
    state_info = text.split("->")[0][:-2].strip('\n')
    num_states=len((state_info.split(";")))-1

    # The last part of the file - contains all of the edge info
    path_info = text.split("->",1)[0]

    # Delete all of the state_info except for the first from_node, edge text, then add the path info back
    path_info = path_info[(len(path_info)-2):].strip('\n') +"->" + (text.split("->",1)[1])

    # declare and init the array - num_states x num_states is probably overkill. Init with -1 since we know that will never be a valid value
    node_paths = [ [-1]*num_states for i in range(num_states)]

    # length * 3 since it splits into 3 parts: first part is the from_node
    # second part is the to_node
    # third part is the edge label. Delete the trailing '}' and new line character. i+=3 to jump to the next line
    for i in range(0,(3*len(path_info.split(";")))-3,3):
    
        # iterate through node_paths[0] to node_paths[num_states] and append the to node and label each time
        # Since it could be out of order, it can't be 0 through num_states - each line you have to check what the from_node actually is, and access that node_paths[#]
        node_paths[int(path_info.replace('->',' ').replace(';',' ').split()[i])].append((path_info.replace('->',' ').replace(';',' ').split()[i+1])+' ' +(path_info.replace('->',' ').replace(';',' ').split()[i+2]))

    # Remove all the initialized -1's from the array
    for i in range(num_states):
        while -1 in node_paths[i]:
            node_paths[i].remove(-1)

    #node_paths[x] contains the nodes that x goes to, plus the edge label.

    curr_node=crits
    next_nodes = []
    parsed_nodes = []
    remove_nodes = []

    with open("temp.dot", 'w') as tmpDOT:
        tmpDOT.write(state_info+"\n")
        while True:
            for i in range(num_states):
                #if curr_node is the to_node, log the from_node
            # if (re.sub(r'^.*?', curr_node+" ", curr_node+" ")) in repr(node_paths[i]):
                if (str(curr_node) + " ") in repr(node_paths[i]):
                    from_node=i
                    index =[j for j, s in enumerate(node_paths[i]) if (str(curr_node)+" ") in s]
                    if (node_paths[i][index[0]].index(str(curr_node) + " ")) != 0:
                        continue
                    if from_node not in next_nodes:
                        next_nodes.append(from_node)
                    tmpDOT.write(str(from_node)+"->"+node_paths[i][index[0]]+";"+"\n")
            if len(next_nodes)>0:
                next_nodes.sort()
                parsed_nodes.append(str(curr_node))
                curr_node=next_nodes.pop()
            else:
                break
        tmpDOT.write("\n")
        parsed_nodes.append('0')
    tmpDOT.close()

    # Delete the unused states from the dot file so there aren't floating nodes
    for i in range(num_states):
        if repr(i) not in parsed_nodes:
            remove_nodes.append(i)

    # Count -1 for the new line char, written to make sure there's not a leading new line in the file
    count = -1
    written=0

    pwd=os.getcwd()
    # Remove all duplicate lines
    #with open(args.output_d, 'w') as outDOT:
    with io.open(pwd +'/path_walk/' + "node_" + str(crits)+".dot", 'w', encoding='utf-8') as outDOT:
        lines_seen = set()
        for line in open("temp.dot", 'r'):
            if count not in remove_nodes:
                if line not in lines_seen:
                    if written == 0:
                        outDOT.write(line.strip('\n'))
                        written=1
                    else:
                        outDOT.write('\n'+ line.strip('\n'))
                        lines_seen.add(line)
            count+=1
        outDOT.write("}")
    outDOT.close()

