#!/usr/bin/env python

import os, sys

import pexpect

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

SAGE        = "%s/sage"%SAGE_ROOT

CLOSE       = '\\end{verbatim}'
OPEN        = '\\begin{verbatim}'
PYTHON      = SAGE_ROOT + "/local/bin/python"
TESTDIR     = ".test"
TIMEOUT     = 10
EXCLUDES    = ["Wall time:", "CPU times:", "may vary"]

SKIP        = "%skip"
SESSION     = "%session"
FULLSESSION = "%fullsession"

ERRORS = ''

def extract_latex_doc(file_name):
    F = open(file_name)
    sessions = []
    n = 0
    prev_line = ''
    in_session = False
    for L in F.xreadlines():
        L = L[:-1]
        if in_session:
            if L == CLOSE:
                in_session = False
                sessions.append(cur)
            else:
                cur.append((n, L))
        elif n > 0 and prev_line.find(SKIP) != -1:
            pass
        elif L == OPEN:
            in_session = True
            cur = []
        n += 1
        prev_line = L
    return sessions

def extract_python_doc(file_name):
    F = open(file_name)
    sessions = []
    n = 0
    prev_line = ''
    in_session = False
    for L in F.xreadlines():
        L = L[:-1]
        if in_session:
            if L.lstrip().rstrip()=='' or L.find('"""') != -1:
                in_session = False
                sessions.append(cur)
            else:
                cur.append((n, L))
        elif n > 0 and prev_line.find(SKIP) != -1:
            pass
        elif len(L) > 0 and L.find('sage:') != -1 and L.lstrip()[0] != '#':
            in_session = True
            cur = [(n,L)]
        n += 1
        prev_line = L
    return sessions

E = None
def initialize_sage():
    global E
    if E is None:
        pass
        #print "Starting SAGE"
    else:
        print "Restarting SAGE"
    E = pexpect.spawn(SAGE)
    E.timeout = TIMEOUT
    E.setmaxread = 10000  # faster
    E.expect('sage:')
    for f in sys.argv[1:]:
        if f[-5:] == '.sage':
            sendline('load "%s"'%f)
            
    #sendline("%Pprint")
    #E.readline()
    #E.readline()

IN = []
def sendline(L):
    global IN
    IN.append(L)
    E.sendline(L)
    E.expect('sage:')

OUT = []
def readlines():
    global OUT
    r = E.before
    j = r.find(chr(10))
    if j != -1:
        r = r[j+1:]
    s = ''
    num_lines = 0
    for L in r.split('\r\n'):
        L = no_escapes(L)
        i = L.find("= ")
        if i != -1:
            L = L[i+2:]
        OUT.append(L)
        s += '\n' + L
        num_lines += 1
    return s, num_lines

def context():
    print "IN:"
    print '\n'.join(IN[-5:])
    print "OUT:"
    print '\n'.join(OUT[-5:])

def no_whitespace(x):
    x = x.replace('<BLANKLINE>','')
    s = ''.join(x.split())
    return s

def output_is_different(x, y):
    if y.find("...") != -1:
        return False
    y = y.lstrip().rstrip()
    x = x.lstrip().rstrip()
    if x[:20] == '-'*20:
        x = x.lstrip('-')[1:].rstrip()
    if x[:9] == 'Traceback' and y[:9] == 'Traceback':
        return no_whitespace(x.split('\n')[-1]) != no_whitespace(y.split('\n')[-1].lstrip())
    x = no_whitespace(x)
    y = no_whitespace(y)
    if x != y:
        print "\nx: len=%s, '%s'"%(len(x), x)
        print "y: len=%s, '%s'"%(len(y), y)
        if len(x) <= len(y):
            for i in range(len(x)):
                if x[i] != y[i]:
                    print "character %s differs: %s, %s"%(i, x[i], y[i])
                    break
    return x != y

def no_escapes(x):
    while True:
        j = x.find(chr(27))
        if j == -1:
            return x
        k = j + x[j:].find(chr(109))
        x = x[:j] + x[k+1:]

def test_session(S, verbose):
    global E
    
    if verbose:
        print "-"*60
        print "\n".join([line for _, line in S])
        print "-"*60

    i = 0
    while i < len(S):
        n, L = S[i]
        if L == "":
            i += 1
            continue
        if L.find("sage:") != -1 or L.find(">>>") != -1:
            while True:
                j = L.find("sage:")
                if j != -1:
                    L = L[j+5:]
                else:
                    break
            sendline(L)
            i += 1
        else:
            M, num_lines = readlines()
            while i+1 < len(S) and S[i+1][1].find('sage:') == -1:
                i += 1
                _, L2 = S[i]
                L += '\n' + L2    
            ignore_diffs =  True in [L.find(exc) != -1 for exc in EXCLUDES]
            if not ignore_diffs and output_is_different(M, L):
                print "** There is an error close to line %s **"%n
                print "\nEXPECTED:\n%s"%L
                print "\nRECEIVED:%s"%M
                context()
                raise RuntimeError
                return
            i += 1
            
            
def test_file(file, start, stop, verbose):
    global ERRORS
    name = os.path.basename(file)
    name = name[:name.find(".")]
    if file[-4:] == '.tex':
        sessions = extract_latex_doc(file)
    elif file[-3:] == '.py' or file[-4:] == '.pyx' or file[-5:]=='.sage':

        sessions = extract_python_doc(file)
    else:
        print "Skipping %s; unknown file type"%file
        return 0
    i = 0
    for S in sessions:
        sys.stdout.flush()
        try:
            if i >= start:
                print "Example %s (line %s)"%(i,S[0][0])                        
                test_session(S, verbose=verbose)
        except pexpect.TIMEOUT:
            print "Example %s (line %s)"%(i,S[0][0])                        
            print "TIMEOUT!!"
            context()
            initialize_sage()
        except KeyboardInterrupt:
            print "Example %s (line %s)"%(i,S[0][0])            
            print "User interrupt."
            return
        except RuntimeError:
            s = "Example %s (line %s)"%(i,S[0][0])
            print s
            print "Documentation contains a mistake.  Please fix and re-run."
            ERRORS += '%s\n'%s
            #return 1
        i += 1
        if stop != 0 and i > stop:
            break
    return 0

if __name__ ==  '__main__':
    import os, sys
    print ''
    if len(sys.argv) == 1:
        print "Usage: %s [--verbose or -v] <filename.tex> [start number] [stop number]"%(sys.argv[0])
        print "Test the documentation in a .tex or .sage file using SAGE (via a pseudo-tty)."
        sys.exit(1)
    else:
        verbose = False
        for i in range(1, len(sys.argv)):
            if sys.argv[i] in ['--verbose', '-verbose', '-v']:
                verbose = True
                del sys.argv[i]
                break
        file = sys.argv[1]
        start = 0
        stop  = 0
        if len(sys.argv) > 2:
            start = int(sys.argv[2])
        if len(sys.argv) > 3:
            stop = int(sys.argv[3])
        initialize_sage()
        if test_file(file, start, stop, verbose=verbose):
            sys.exit(1)
        if len(ERRORS) > 0:
            print ERRORS
            sys.exit(1)
        e = E.terminate(0)
        if e:
            print " ** Unclean exit -- possibly a memory error or segmentation fault occurred. ** "
            sys.exit(1)
        try:
            # Hack to get around stupid pointless error message.
            def foo():
                pass
            E.__del__ = foo
            del E
        except:
            pass
        sys.exit(0)
