#!/usr/bin/env python

# the default timeout for doctests: 6 minutes
TIMEOUT = 360
# the timeout value for long doctests: 30 minutes
TIMEOUT_LONG = 1800
# the timeout for doctests running under valgrind tools: unreasonably long
TIMEOUT_VALGRIND = 1024*1024

import os, re, sys, signal, time, shutil, tempfile

optional = False
long_time = False
verbose = False

argv = sys.argv

import sage.misc.preparser

######################################################
# This code is copied from sage.misc.misc for speed:
######################################################
def is_64bit():
    return sys.maxint == 9223372036854775807

DOT_SAGE = os.environ['DOT_SAGE']
SAGE_TMP='%s/tmp/%s/'%(DOT_SAGE,os.getpid())
def delete_tmpfiles():
    try:
        pass
#        shutil.rmtree(SAGE_TMP)
    except OSError:
        pass

######################################################


SAGE_ROOT=os.environ["SAGE_ROOT"]
LD = os.environ["LD_LIBRARY_PATH"]
os.environ["LD_LIBRARY_PATH"] = SAGE_ROOT + "/local/lib:" + LD
os.environ["PYTHONPATH"]=SAGE_ROOT + "/local/lib/python/site-packages"

PYTHON = SAGE_ROOT + "/local/bin/python"

def test_code():
    return """
if __name__ ==  '__main__':
    import doctest, sys
    s = doctest.testmod(sys.modules[__name__],
                   optionflags=doctest.NORMALIZE_WHITESPACE
                              |doctest.ELLIPSIS,
                   verbose=%s,
                   globs=globals())
    quit_sage(verbose=False)
"""%verbose

ALARM = """
import signal
def __mysig__(a,b):
    print "TIME OUT TIME OUT TIME OUT"
    raise RuntimeError, "timeout"
signal.signal(signal.SIGALRM, __mysig__)
signal.alarm(%s)
"""%TIMEOUT

NONE=0; LONG_TIME=1; RANDOM=2; OPTIONAL=3; NOT_IMPLEMENTED=4; NOT_TESTED=5

def comment_modifier(s):
    sind = s.find('#')
    if sind == -1 or s[sind:sind+3] == '###':
        return []
    eind = s.find('###',sind+1)
    L = s[sind+1:eind].lower()
    v = []
    if ('optional' in L) or ('known bug' in L):
        v.append(OPTIONAL)
    if 'long time' in L:
        v.append(LONG_TIME)
    if 'not implemented' in L:
        v.append(NOT_IMPLEMENTED)
    if 'not tested' in L:
        v.append(NOT_TESTED)
    if 'random' in L:
        v.append(RANDOM)
    return v

def preparse_line_with_prompt(L):
    i = L.find(':')
    if i == -1:
        return sage.misc.preparser.preparse(L)
    else:
        return L[:i+1] + sage.misc.preparser.preparse(L[i+1:])

def doc_preparse(s):
    """
    Run the preparser on the documentation string s.
    This *only* preparses the input lines, i.e., those
    that begin with "sage:".or with "..."
    """
    sl = s.lower()

    if not optional and sl.find("optional") != -1 and \
               sl.find('package') != -1 and sl.find('installed'):
        return ''

    # Deal with code whose output should be ignored. 
    t = []
    for L in s.split('\n'):
        begin = L.lstrip()[:5]
        if begin == 'sage:':
            c = comment_modifier(L)
            line = ''
            if LONG_TIME in c and not long_time:
                L = '\n'  # extra line so output ignored
            if RANDOM in c:
                L = L.replace('sage:', 'sage: print "ignore this"; ') 
            if NOT_TESTED in c:
                L = '\n'   # not tested
            if NOT_IMPLEMENTED in c:
                L = '\n'   # not tested
            if OPTIONAL in c and not optional:
                L = '\n'
            line = preparse_line_with_prompt(L)
            if RANDOM in c:
                # count spaces at the beginning of line to fix alignment later
                i = 0
                while i < len(line) and line[i].isspace():
                    i += 1
                # append a line saying 'ignore' followed by ellipsis (...)
                # and an empty line, to ignore the output given in the test
                line += '\n' + ' '*i + 'ignore ...\n'
            t.append(line)
            
        elif begin.startswith('...'):

            i = L.find('.')
            t.append(L[:i+3] + sage.misc.preparser.preparse(L[i+3:]))
            
        else:
            
            t.append(L)
    return '\n'.join(t)

def extract_doc(file_name, module):
    a = os.path.abspath(file_name)
    i = a.rfind("sage/")
    #module_name = a[i:].replace("/",".")
    #j = module_name.rfind(".")
    #module_name = module_name[:j]

    F = open(file_name).readlines()
    # Put line numbers on every input line
    v = []
    i = 1
    j = 0
    while j < len(F):
        L = F[j].rstrip()
        if L.lstrip()[:5] == 'sage:':
            while j < len(F) and L.endswith('\\') and not L.endswith('\\\\'):
                j += 1
                i += 1
                L = L[:-1] + F[j].lstrip().lstrip('...').rstrip()
            #L += '\t\t#<--- "LINE %s of %s"'%(i, file_name)
            L += '###_sage"line %s:_sage_    %s_sage"'%(i, L[4:].strip())
        j += 1
        i += 1
        v.append(L)

    # 32/64-bit.  If we're on a 32-bit computer, remove all lines that
    # contains "# 64-bit", and if we're on a 64-bit machine remove all
    # lines that contain "# 32-bit".  This makes it possible to have
    # different output in the doctests for different bit machines.
    
    if is_64bit():
        exclude_string = "# 32-bit"
    else:
        exclude_string = "# 64-bit"
        
    F = '\n'.join([L for L in v if not exclude_string in L])

    if is_64bit():
        F = F.replace('# 64-bit','')
    else:
        F = F.replace('# 32-bit','')
    
    F = F.replace('\'"""\'','')

    if file_name[-4:] == ".tex":
        F = pythonify(F)
    
    n = 0
    i = 0
    s = "# -*- coding: utf-8 -*-\n"
    s += "from sage.all_cmdline import *; \n"
    s += "import sage.plot.plot; sage.plot.plot.DOCTEST_MODE=True\n"  # turn off image popup

    while True:
        i = F.find('"""')
        if i == -1: break
        name = "example"        
##        z = F[:i].rfind('def ')
##         if z != -1:
##             z2 = F[z+4:].find(':')
##             z3 = F[z+4:].find('(')
##             z4 = min(z2,z3)
##             if z4 != -1:
##                 try:
##                     name = F[z+4:z+4+z4].strip().split()[0]
##                 except IndexError:
##                     pass
##         if '.' in name:
##             name = 'example'
        k = F[i+3].find('"""')
        j = i+3 + F[i+3:].find('"""')
        s += "def %s_%s():"%(name,n)
        n += 1
        try:
            doc = doc_preparse(F[i:j+3])
        except SyntaxError:
            doc = F[i:j+3]
        if len(doc):
            doc = '""">>> set_random_seed(0L)\n\n' + doc[3:]
        s += "\tr"+ doc + "\n\n"
        F = F[j+3:]

    #i = F.find('def __doctest_cleanup')
    #if i != -1:
    #    s += F[i:] + '\n'
    #    s += test_code()
    #    s += '    __doctest_cleanup()\n'
    #else:
    s += test_code()
    s += '    sys.exit(s[0])'


    # Allow for "sage:" instead of the traditional Python ">>>".
    s = s.replace("sage:",">>>").replace('_sage"','')

    return s

def pythonify(F):
    """
    INPUT:
        F -- string; read in latex file
    OUTPUT:
        string -- python program that has functions with docstrings made from the
                  verbatim examples in the latex file.
    """
    # Close links:
    F = F.replace('\\end{verbatim}%link','')
    F = F.replace('%link\n\\begin{verbatim}','')
    
    # Get rid of skipped code
    s = ''
    while True:
        i = F.find('%skip')
        if i == -1:
            s += F
            break
        s += F[:i]
        F = F[i:]
        j = F.find('\\end{verbatim}')
        if j == -1:
            break
        F = F[j + len('\\end{verbatim}')+1:]
    F = s
    
    # Make the verbatim environ's get extracted via the usual parser above
    F = F.replace("\\begin{verbatim}",'"""')
    F = F.replace("\\end{verbatim}",'"""')
    return F

def post_process(s, file, tmpname):
    s = s.replace('.doctest_','')
    if file[-3:] == 'pyx':
        s = s.replace('"%s", line'%file[:-1], '"%s", line'%file)
    i = s.find("Failed example:")
    cnt = 0
    while i != -1:
        k = s[:i].rfind('File')
        k += s[k:].find(',')
        j = s[i:].find('###line')
        s = s[:k] + ', ' + s[i+j+3:]
        i = s.find("Failed example:")
        cnt += 1
        if cnt > 1000:
            break
    s = s.replace(':_sage_',':\n').replace('>>>','sage:')
    c = '###line [0-9]*\n'
    r = re.compile(c)
    s = r.sub('\n',s)
    if cnt > 0:
        s += "For whitespace errors, see the file %s"%tmpname
    return s

__alarm_time=0
def __mysig(a,b):
    raise RuntimeError, "computation timed out because alarm was set for %s seconds"%__alarm_time

def alarm(seconds):
    """
    Raise a KeyboardInterrupt exception in a given number of seconds.
    This is useful for automatically interrupting long computations
    and can be trapped using exception handling.

    INPUT:
        seconds -- integer
    """
    seconds = int(seconds)
    # Set our alarm signal handler. 
    signal.signal(signal.SIGALRM, __mysig)
    global __alarm_time
    __alarm_time = seconds
    signal.alarm(seconds)

def test_file(file):
    if os.path.exists(file):
        name = os.path.basename(file)
        name = name[:name.find(".")]
        s = extract_doc(file, name)
        if 'SAGE_TESTDIR' not in os.environ or os.environ['SAGE_TESTDIR'] is "":
            os.environ['SAGE_TESTDIR'] = os.environ['SAGE_ROOT']+"/tmp"
        f = "%s/.doctest_%s.py"%(os.environ['SAGE_TESTDIR'],name)
        d = "%s/.doctest"%os.environ['SAGE_TESTDIR']
        if not os.path.exists(d):
            os.makedirs(d)
        open(f,"w").write(s)
        cmd = "%s %s"%(PYTHON, f)
        if gdb:
            print "*"*80
            print "Type r at the (gdb) prompt to run the doctests."
            print "Type bt if there is a crash to see a traceback."
            print "*"*80
            cmd = "gdb --args " + cmd

        if memcheck:
            cmd= "valgrind --tool=memcheck --trace-children=yes --leak-resolution=high --log-file=$HOME/.sage/valgrind/sage-memcheck.%p --show-reachable=yes --leak-check=full " + cmd
        if massif:
            cmd = "valgrind --tool=massif --trace-children=yes --depth=6 --log-file=$HOME/.sage/valgrind/sage-massif.%p " + cmd
        if cachegrind:
            cmd = "valgrind --tool=cachegrind --trace-children=yes --log-file=$HOME/.sage/valgrind/sage-cachegrind.%p " + cmd
        if omega:
            cmd = "valgrind --tool=exp-omega --trace-children=yes --log-file=$HOME/.sage/valgrind/sage-omega.%p " + cmd

        VALGRIND = '%s/valgrind/'%DOT_SAGE
        if not os.path.exists(VALGRIND):
          os.makedirs(VALGRIND)

        tm = time.time()
        try:
            if verbose or gdb or memcheck or massif or cachegrind:
                e = os.system(cmd)
                out = ''; err = ''
            else:
                outf = tempfile.NamedTemporaryFile()
                errf = tempfile.NamedTemporaryFile()
                e = os.system("%s 1>%s 2>%s"%(cmd,outf.name,errf.name))
                out = outf.read()
                err = errf.read()
        except KeyboardInterrupt:
            err += "\n Error -- interrupted after %s seconds!"%tm
            print "Error!!!"
            sys.exit(2)

        if not verbose and 'raise KeyboardInterrupt' in err:
            print "\n" + "*"*80 + "\nControl-C pressed -- interrupting doctests.\n" + "*"*80
            sys.exit(2)

        if time.time() - tm >= TIMEOUT:
            err += "*** *** Error: TIMED OUT! *** ***"
            print err
            
        s = post_process(out, file, f)  + err
        print s
        numfail = s.count('Failed') + s.count('Error')
        if numfail > 0:
            delete_tmpfiles()
            sys.exit(4)
        elif e != 0:
            if not verbose:
                print "A mysterious error (perphaps a memory error?) occurred, which may have crashed doctest."
            sys.exit(3)
        else:
            delete_tmpfiles()
            sys.exit(0)
    else:
        print "Error running %s, since file %s does not exist."%(
            argv[0], argv[1])
        delete_tmpfiles()
        sys.exit(1)


def has_opt(opt):
    if '-' + opt in argv:
        i = argv.index('-' + opt)
        del argv[i]
        return True
    elif '--' + opt in argv:
        i = argv.index('--' + opt)
        del argv[i]
        return True
    return False

def usage():
    print "\n\nUsage: sage -t [-optional] [-long] [-verbose] <filename.py | filename.pyx>"
    print "If -optional is present include examples that uses optional packages."
    print "If -long is present include examples that take a long time to run,"
    print 'which is defined by the input line containing the phrase "long time".'
    sys.exit(1)
        
if __name__ ==  '__main__':
    import os, sys
    if len(argv) == 1:
        usage()
    else:
        if has_opt('help') or has_opt('h') or has_opt('?'):
            usage()
        optional   = has_opt('optional')
        long_time  = has_opt('long')
        verbose    = has_opt('verbose')
        gdb        = has_opt('gdb')
        memcheck   = has_opt('memcheck') or has_opt('valgrind')
        massif     = has_opt('massif')
        cachegrind = has_opt('cachegrind')
        omega      = has_opt('omega')
        if long_time:
            TIMEOUT = TIMEOUT_LONG
        if gdb or memcheck or massif or cachegrind or omega:
            TIMEOUT = TIMEOUT_VALGRIND
        if argv[1][0] == '-':
            usage()
        try:
            test_file(argv[1])
        except KeyboardInterrupt:
            sys.exit(1)
            
        

