#!/usr/bin/env python3

'''
 * File: listBehaviors.py
 *
 * Author: Brad Neuman
 * Created: 2018-05-07
 *
 * Description: Simple script to vizualize output from the DelegationTree.DumpBehaviorTreeBranchesToFile unit test
 *               in an easier to read text format
 *
 * Copyright: Anki, Inc. 2018
'''

import argparse
import os


def parseBranchesFile(infile):
    '''
    given a behavior_branches.txt file (from engine unit test output), return a nested dictionary representing
    the behavior tree and any extra csv data associated
    '''
    
    ret = {}
    
    with open(infile, 'r') as f:
        for line in [x.strip() for x in f.readlines()]:
            tokens = [x.strip() for x in line.split(',')]
            stack = tokens[0]
            tokens = [t for t in tokens[1:] if t]

            stack = stack.split('/')
                
            node = ret;
            for s in stack:
                if not s in node:
                    node[s] = ({}, tokens)
                node = node[s][0]

    return ret

def flattenTree( tree, indent = 0 ):
    ''' turn tree into a flat list of (indent, string) pairs suitable for display '''

    ret = []
    
    for key in tree:
        ret.append( (indent, key, tree[key][1]) );
        ret.extend( flattenTree( tree[key][0], indent+1 ) );

    return ret

def printFlatTree( tree ):
    '''
    given a flat tree (list of (indent : int, behavior : string) pairs), pretty-print an ascii
    representation of the tree with some visual aids
    '''

    # if this many rows exist between entries at a given depth, draw a vertical line to help see which things
    # are at which levels
    MIN_ROWS_FOR_CONNECTOR = 5

    strs = []

    # map of indent depth -> idx of the last entry at this depth or lower
    lastRow = {}

    lastIndent = -1
    
    for idx in range(len(tree)):
        row = tree[idx]
        indent = row[0]
        extraTokenStr = '' if not row[2] else '(' + ', '.join(row[2]) + ')'
        strs.append(list('{}* {} {}'.format('  '*indent, row[1], extraTokenStr)))

        if indent in lastRow:
            if idx - lastRow[indent] >= MIN_ROWS_FOR_CONNECTOR:
                # fill the line now to everything between here and there
                for i2 in range(lastRow[indent]+1, idx):
                    if len(strs[i2]) > indent*2:
                        strs[i2][ indent * 2 ] = '|'
                    else:
                        raise ValueError(f'invalid modification: indent={indent}, str={strs[i2]}')

        lastRow[indent] = idx

        # also reset anything deeper than indent
        for iKey in lastRow:
            if iKey > indent:
                lastRow[iKey] = idx

    for s in strs:
        print(''.join(s))
    

def main():
    
    parser = argparse.ArgumentParser(description='Prettier text output for behavior_branches.txt')
    parser.add_argument('inputFilename', metavar='INPUTFILE', nargs='?',
                        default='./_build/mac/Debug/test/engine/behavior_branches.txt',
                        help='the file generated by the unit test that dumps all behavior branches')

    args = parser.parse_args()
    inputFilename = args.inputFilename

    try:
        behaviorTree = parseBranchesFile(inputFilename)
    except FileNotFoundError:
        print('file {} not found. Did you run the engine unit test DelegationTree.DumpBehaviorTreeBranchesToFile ?'.\
              format(inputFilename))
        return

    flatTree = flattenTree(behaviorTree)

    printFlatTree( flatTree )

main()
