#! /usr/bin/env python
#
# Extract test programs from the Chapel spec
#
# All tests are place in a subdirectory of the output directory
# (default '.') named after the LaTeX file.  Multiple LaTex files may
# be specified on the command line.
#
# See spec/chapel_testing.tex and the LaTex source for the spec itself
#  for more info.
#
# SUPPORTED ENVIROMENTS:
# - chapelexample (REQUIRED)
#   If a filename is not given on the first line, the filename is
#   generated from the input filename and the line number.
#
# - chapelpre
# - chapel|chapelcode|chapelnoprint
# - chapelpost
# - chapelfuture
# - chapelcompopts
# - chapelexecopts
# - chapelprediff
# - chapeloutput|chapelprintoutput
# - chapelwideoutput
#
# GENERAL CAVEATS:
# - All environments may contain special formatting due the fact that
#   everything must be able to be parsed by LaTeX.  Currently, the
#   following substitutions are made:
#
#      '//' becomes '/'
#      '/#' becomes '#'
#      '/_' becomes '_'
#      '/&' becomes '&'
#      '$\mbox{{\bf \$}}$' becomes '$'
#      '(*\texttt{\$}*)' becomes '$'
#
#   If this becomes too unwieldy, we'll probably have to find a better
#   way to do it.
#

import os, sys

#
# Helper functions
#

# Extract the good file name from a source line
def GetGoodName(line):
    outfile = line.strip().strip('%').strip().replace('.good', '')
    if options.debug:
        sys.stdout.write('>>> %s\n'%(outfile))
    currTestNames.append(outfile)
    if currTestNames.count(outfile) > 1:
        sys.stdout.write("WARNING: Test name '%s' already specified.  Good filename may conflict\n"%(outfile))
    return outdir+'/'+curdir+'/'+outfile+'.good'

# Helper functions for de-LaTeX-ifying output line
def RemoveLaTeXFormatting(line, nlrepl):
    dline = line.replace('\\#', '#')
    dline = dline.replace('\\&', '&')
    dline = dline.replace('\\_', '_')
    dline = dline.replace('\\\\', nlrepl)
    dline = dline.replace('$\\mbox{{\\bf \\$}}$', '$')
    dline = dline.replace('(*\\texttt{\\$}*)', '$')
    return dline

# Helper for writing test files
def WriteTestFile(fh, output):
    for line in output:
        if line[0] != '%':
            fh.write(RemoveLaTeXFormatting(line, '\\'))

# Helper for writing test .good files
def WriteGoodFile(fh, output, testName):
    goodTestName = testName+'.good'
    opened = False
    for line in output:
        if line[0] == '%':
            # new .good file
            if opened:
                fh.close()
            opened = True
            goodTestName = GetGoodName(line)
            if options.debug:
                sys.stdout.write('>>> Opening %s\n'%(goodTestName))
            fh = open(goodTestName, 'w')
        else:
            if not opened:
                if options.debug:
                    sys.stdout.write('>>> Opening %s\n'%(goodTestName))
                fh = open(goodTestName, 'w')
                opened = True
            fh.write(RemoveLaTeXFormatting(line, ''))

    if opened:
        fh.close()


def ClearLists():
    del chapelpre[:]
    del chapelcode[:]
    del chapelpost[:]
    del chapelfuture[:]
    del chapelcompopts[:]
    del chapelexecopts[:]
    del chapelprediff[:]
    del chapeloutput[:]
    del chapelwideoutput[:]
    del chapelprintoutput[:]

# Write out the test files
def OutputTest(lineno, currTestName):
    global TotalTests, TotalFutures, TotalIgnored
    if (len(chapelfuture) > 0) and not options.gen_futures:
        TotalIgnored += 1
        if options.verbose:
            sys.stdout.write('\tIgnoring future %s.chpl on line %d\n'%(currTestName, lineno))
        return;

    TotalTests += 1
    # writeout the actual test program
    if (len(currTestName)-currTestName.rfind('.chpl')) == 5:
        testName = currTestName[:currTestName.rfind('.chpl')]
    else:
        testName = currTestName

    currTestNames.append(testName)
    if currTestNames.count(testName) > 1:
        sys.stdout.write("WARNING: Test name '%s' already specified\n"%
                         (testName))
        newof = testName + '.' + str(currTestNames.count(testName))
        currTestNames.append(newof)
        sys.stdout.write('         Renaming it to %s\n'%(newof));
        testName = newof

    # add full path
    testName = outdir+'/'+curdir+'/'+testName

    if options.verbose:
        sys.stdout.write('\tWriting %s.chpl (line %d)\n'
                         %(testName.rsplit('/', 1)[1], lineno))

    fh = open(testName+'.chpl', 'w')
    if options.debug:
        sys.stdout.write('>>> Opening %s\n'%(testName+'.chpl'))

    WriteTestFile(fh, chapelpre)
    WriteTestFile(fh, chapelcode)
    WriteTestFile(fh, chapelpost)

    fh.close()

    # write out the future file
    if len(chapelfuture) > 0:
        TotalFutures += 1
        if options.debug:
            sys.stdout.write('>>> Opening %s\n'%(testName+'.future'))
        fh = open(testName+'.future', 'w')
        WriteTestFile(fh, chapelfuture)

        # write out the .compopts file
    if len(chapelcompopts) > 0:
        if options.debug:
            sys.stdout.write('>>> Opening %s\n'%(testName+'.compopts'))
        fh = open(testName+'.compopts', 'w')
        WriteTestFile(fh, chapelcompopts)

    # write out the .execopts files
    if len(chapelexecopts) > 0:
        if options.debug:
            sys.stdout.write('>>> Opening %s\n'%(testName+'.execopts'))
        fh = open(testName+'.execopts', 'w')
        WriteTestFile(fh, chapelexecopts)

    # write out the .prediff file
    if len(chapelprediff) > 0:
        if options.debug:
            sys.stdout.write('>>> Opening %s\n'%(testName+'.prediff'))
        fh = open(testName+'.prediff', 'w')
        WriteTestFile(fh, chapelprediff)
        os.chmod(testName+'.prediff', 0o755)

    # write out .good files
    if len(chapeloutput) > 0 or len(chapelprintoutput) > 0:
        WriteGoodFile(fh, chapeloutput, testName)
        WriteGoodFile(fh, chapelprintoutput, testName)
    else:
        # write an empty .good file
        fh = open(testName+'.good', 'w')
        fh.close()

    # write out .no-local.good file if it exists.
    if len(chapelwideoutput) > 0:
        WriteGoodFile(fh, chapelwideoutput, testName+'.no-local')
        WriteGoodFile(fh, chapelprintoutput, testName+'.no-local')

    ClearLists()





####################
# extract_tests.py #
####################

#
# Initialize global vars
#
chapelpre = list()
chapelcode = list()
chapelpost = list()
chapelfuture = list()
chapelcompopts = list()
chapelexecopts = list()
chapelprediff = list()
chapeloutput = list()
chapelwideoutput = list()
chapelprintoutput = list()
current = None
currTestNames = list()

#
# Accounting
#
TotalTests = 0
TotalExamples = 0
TotalFutures = 0
TotalIgnored = 0

#
# Parse command line args
#
from optparse import OptionParser

parser = OptionParser("usage: %prog [options] file1.tex file2.tex ...")
parser.add_option('-o', '--outdir',
                  action='store', type='string', dest='outdir',
                  default='.', help="output directory (default '.')",
                  metavar='DIR')
parser.add_option('', '--no-futures',
                  action='store_false', dest='gen_futures', default=True,
                  help="don't generate future tests")
parser.add_option('-q', '--quiet',
                  action='store_false', dest='verbose', default=True,
                  help="don't print status messages to stdout")
parser.add_option('-d', '--debug',
                  action='store_true', dest='debug', default=False)

(options, args) = parser.parse_args()

if len(args) < 1:
    parser.error("incorrect number of arguments")

#
# Create output directory
#
outdir = os.path.realpath(options.outdir)
if not os.path.isdir(outdir):
    os.mkdir(outdir)

currTestName = None

for infile in args:
    #
    # Read the input file
    #
    realinfile = os.path.realpath(infile)
    myfile = open(realinfile, 'r')
    myLines = myfile.readlines()
    myfile.close()

    # Clean up table of testname
    del currTestNames[:]

    curinfile = realinfile
    # curdir = realinfile.rsplit('/, 1)[1]+'.tests'
    curdir = (realinfile.rsplit('/', 1)[1]).rsplit('.', 1)[0]
    if not os.path.isdir(outdir+'/'+curdir):
        os.mkdir(outdir+'/'+curdir)
    if options.verbose:
        sys.stdout.write('***\n')
        sys.stdout.write('*** Generating tests from from %s\n'%(infile))
        sys.stdout.write('***\tOutput directory: %s/%s\n'%
                         (options.outdir,curdir))

    lineno = 0
    testlineno = lineno
    printToOutfile = False

    for line in myLines:
        lineno += 1
        pos = line.find('\\end{chapel')
        if (line[0].strip() != '%') and pos > -1:
            pos += 11
            if line.find('example}', pos) == pos:
                if options.debug:
                    sys.stdout.write('>>> Writing out test example: %s\n'%(currTestName))
                OutputTest(testlineno, currTestName)
                currTestName = None
                ClearLists()
                continue

            printToOutfile = False
            current = None
            continue

        if printToOutfile:
            current += [line]
            continue

        pos = line.find('\\begin{chapel')
        if (line[0].strip() != '%') and pos > -1:
            pos += 13
            if options.debug:
                sys.stdout.write('>>> %s\n'%(line))
            if (line.find('example}{', pos) == pos):
                testlineno = lineno
                currTestName = line[pos+9:-2]
                current = None
                if options.debug:
                    sys.stdout.write('>>> Found Chapel test example: %s\n'%(currTestName))
                continue

            if currTestName == None:
                if ((line.find('}', pos) == pos) or
                    (line.find('code}', pos) == pos)):
                    TotalIgnored += 1
                    if options.verbose:
                        sys.stdout.write('\tIgnoring Chapel code on line %d\n'%(lineno))
                continue

            if (line.find('pre}', pos) == pos):
                printToOutfile = True
                current = chapelpre;

            elif ((line.find('}', pos) == pos) or
                  (line.find('code}', pos) == pos) or
                  (line.find('noprint}', pos) == pos)):
                printToOutfile = True
                current = chapelcode;

            elif line.find('post}', pos) == pos:
                printToOutfile = True
                current = chapelpost;

            elif line.find('future}', pos) == pos:
                printToOutfile = True
                current = chapelfuture;

            elif line.find('compopts}', pos) == pos:
                printToOutfile = True
                current = chapelcompopts;
            elif line.find('execopts}', pos) == pos:
                printToOutfile = True
                current = chapelexecopts;
            elif line.find('prediff}', pos) == pos:
                printToOutfile = True
                current = chapelprediff;

            elif line.find('output}', pos) == pos:
                printToOutfile = True
                current = chapeloutput;
            elif line.find('wideoutput}', pos) == pos:
                printToOutfile = True
                current = chapelwideoutput;
            elif line.find('printoutput}', pos) == pos:
                printToOutfile = True
                current = chapelprintoutput;
                pos += 12
                # Extract an explicit good file name
                if line.find('{', pos) == pos:
                    goodName = line[pos+1:-2]
                    # Drop the parameter into the output stream as a comment
                    # only if non-null.
                    if len(goodName) > 0:
                        current += [' '.join(('%',goodName,'\n'))]

        pos = line.find('\\begin{example}')
        if (line[0].strip() != '%') and pos > -1:
            TotalExamples += 1

    if not os.listdir(outdir+'/'+curdir):
        os.rmdir(outdir+'/'+curdir)
        if options.verbose:
            sys.stdout.write('*** Removing empty directory %s/%s\n'%
                             (options.outdir,curdir))


if options.verbose:
    sys.stdout.write("DONE\n\n");

sys.stdout.write('Total tests generated: %d\n'%(TotalTests))
sys.stdout.write('\tFutures: %d\n'%(TotalFutures))
sys.stdout.write('Ignored examples: %d\n'%(TotalExamples))
sys.stdout.write('Ignored code segments: %d\n'%(TotalIgnored-TotalExamples))
