#!/usr/bin/env python

# Overview
# --------
# This is a centralized script for running VTR.
# 
# It wraps various other scripts for running VTR and parses the results
# into sub-commands.

# This provides a straight forward interface based on sub-commands
# (like 'git commit').
# 
# Run as 'vtr help' for usage.

# Adding Sub-commands
# ----------------------
# Each sub-command is a stand-alone python script which is loaded by
# this script at run-time.
#
# To add a new sub-command simply create a file named:
#       
#       vtr-{command}.py
#
# located under:
#    
#       $VTR_ROOT/vtr_flow/scripts
#
# where {command} is the command name.
#
# For example the 'flow' command is implemented in the file:
#
#       $VTR_ROOT/vtr_flow/scripts/vtr-flow.py
#
# Each command script is expected to implement the following two methods:
#
#       def vtr_command_argparser(prog=None):
#           #Setup the argument parser
#           #... 
#           return argparser

# and
#
#       def vtr_command_main(arg_list, prog=None):
#           args = vtr_command_argparser(prog).parse_args(arg_list)
#
#           #Do something useful, based on args
#           #...
#
# vtr_command_argparser() takes the program name as it's first argument and returns
# the argparse.ArgumentParser used by the command (this is used to implement the 'help'
# subcommand).
#
# vtr_command_main() takes the list of command line arguments and the program name
# and executes the command


import argparse
import sys
import os
import textwrap
import importlib
from fnmatch import fnmatch
from collections import OrderedDict

def determine_vtr_root():
    """
    Returns the root of the VTR source tree as determined by:
      1. The VTR_ROOT environment variable, if set otherwise
      2. The location of this script file
    """
    if 'VTR_ROOT' in os.environ:
        return os.environ['VTR_ROOT']
    else:
        #Infer from the directory containing this script file
        return os.path.dirname(os.path.abspath(__file__))

def import_vtr_commands(vtr_root):
    """
    Loads the vtr-*.py scripts from $VTR_ROOT/vtr_flow/scripts as modules 
    Returns a dictionary mapping from the command name to it's module
    """
    unordered_commands = {}

    vtr_command_dir = os.path.join(vtr_root, "vtr_flow", "scripts")
    sys.path.insert(0, vtr_command_dir)

    for file in os.listdir(vtr_command_dir):
        if fnmatch(file, "vtr-*.py"):
            command_base, ext = os.path.splitext(file)

            command_name = command_base.split('-', 1)[1]

            unordered_commands[command_name] = importlib.import_module(command_base)

    #Re-build the command dictionary in preferred order, this ensures
    #the preferred (i.e. common) commands appear first in the help
    commands = OrderedDict()
    for command_name in ["flow", "task"]:
        if command_name in unordered_commands:
            commands[command_name] = unordered_commands[command_name]
            del unordered_commands[command_name]

    for command_name, module in unordered_commands.iteritems():
        commands[command_name] = module

    return commands

VTR_ROOT = determine_vtr_root()
sys.path.insert(0, os.path.join(VTR_ROOT, 'vtr_flow', 'scripts', 'python_libs')) #For verilogtorouting module

from verilogtorouting.util import RawDefaultHelpFormatter

VTR_COMMANDS = import_vtr_commands(VTR_ROOT)

def argument_parser(vtr_commands):

    main_description = textwrap.dedent(
                    """
                    Script for running the Verilog-to-Routing design flow, 
                    conducting experiments and collecting results.
                    """
                  )

    main_epilog = textwrap.dedent(
                    """
                    Examples
                    --------

                        Get help with the 'flow' subcommand:

                            %(prog)s help flow

                        Run the standard design flow for a single architecture and circuit

                            %(prog)s flow arch.xml circuit.v

                        Run a VTR task named 'my_task'

                            %(prog)s task my_task

                    """
              )

    parser = argparse.ArgumentParser(description=main_description,
                                     epilog=main_epilog,
                                     formatter_class=RawDefaultHelpFormatter
                                     )

    subparsers = parser.add_subparsers(title='VTR subcommands',
                                       #help="VTR subcommands",
                                       dest="command")

    help_description = textwrap.dedent(
                    """
                    Get help with subcommands
                    """
                  )

    help_epilog = textwrap.dedent(
                    """
                    Examples
                    --------

                        Get help with the 'flow' subcommand:

                            %(prog)s flow

                        Get help with the 'task' subcommand:

                            %(prog)s task

                    """
              )
    help_parser = subparsers.add_parser('help', 
                                        description=help_description,
                                        help=help_description,
                                        epilog=help_epilog,
                                        formatter_class=RawDefaultHelpFormatter)
    help_parser.add_argument("sub_command",
                             nargs="?",
                             help="Sub-command to get help with")

    for command, module in vtr_commands.iteritems():

        if not hasattr(module, "vtr_command_argparser"):
            raise AttributeError("VTR command '{}' is missing required function '{}'".format(command, "vtr_command_argparser"))
        else:
            description = module.vtr_command_argparser().description

        subparsers.add_parser(command, 
                              add_help=False, #Show full usage if -h supplied
                              help=module.vtr_command_argparser().description)

    return parser, help_parser

def main():
    parser, help_parser = argument_parser(VTR_COMMANDS)

    #Parse the subcommand. Any unkown arguments are passed to the subcommand.
    args, unkown_args = parser.parse_known_args()

    #The effective program being executed (used to make help messages consistent)
    prog = "{} {}".format(os.path.basename(sys.argv[0]), args.command)

    if args.command == "help":

        if not args.sub_command:
            #This script's help
            parser.print_help()

        elif args.sub_command == "help":
            help_parser.print_help()

        elif args.sub_command in VTR_COMMANDS:
            #Show the sub-command in the help output
            help_prog = "{} {}".format(os.path.basename(sys.argv[0]), args.sub_command)

            if not hasattr(VTR_COMMANDS[args.sub_command], "vtr_command_argparser"):
                raise AttributeError("VTR command '{}' is missing required function '{}'".format(args.sub_command, "vtr_command_argparser"))
            else:
                cmd_parser = VTR_COMMANDS[args.sub_command].vtr_command_argparser(help_prog)

            cmd_parser.print_help()

        else:
            print "Unrecognized help subcommand {}".format(args.sub_command)
            parser.print_help()
            sys.exit(1)

    elif args.command in VTR_COMMANDS:
        #Run the appropraite command
        if not hasattr(VTR_COMMANDS[args.command], "vtr_command_main"):
            raise AttributeError("VTR command '{}' is missing required function '{}'".format(args.command, vtr_command_main))
        else:
            VTR_COMMANDS[args.command].vtr_command_main(unkown_args, prog=prog)

    else:
        print "Unrecognized subcommand {}".format(args.help_for_command)
        parser.print_help()
        sys.exit(1)

    sys.exit(0)


if __name__ == "__main__":
    main()