#!/usr/bin/env python
###############################################################################
#                                                                             #
#   Copyright 2005 University of Cambridge Computer Laboratory.               #
#                                                                             #
#   This file is part of Nprobe.                                              #
#                                                                             #
#   Nprobe is free software; you can redistribute it and/or modify            #
#   it under the terms of the GNU General Public License as published by      #
#   the Free Software Foundation; either version 2 of the License, or         #
#   (at your option) any later version.                                       #
#                                                                             #
#   Nprobe is distributed in the hope that it will be useful,                 #
#   but WITHOUT ANY WARRANTY; without even the implied warranty of            #
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the             #
#   GNU General Public License for more details.                              #
#                                                                             #
#   You should have received a copy of the GNU General Public License         #
#   along with Nprobe; if not, write to the Free Software                     #
#   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA #
#                                                                             #
###############################################################################


##############################################################################
## 
## 
## Search Python files/directories for text, regular expression, or numbered
##  lines, presenting matches in context
##  
##

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

from string import *
from sys import argv
import getopt
import os
import sys
from stat import *
from os.path import basename, dirname, join, isfile, isdir

import popen2
import re 

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


#
# Use pretty printing if possible
#

try:
    from print_col import cprint, cbprint, bprint
    from print_col import F_RED, F_GREEN, F_BLUE, F_MAGENTA, F_CYAN, F_YELLOW
    def print_red(s):
        cprint(F_RED, s)
    def print_green(s):
        cprint(F_GREEN, s)
    def print_blue(s):
        cprint(F_BLUE, s)
    def print_magenta(s):
        cprint(F_MAGENTA, s)
    def print_cyan(s):
        cprint(F_CYAN, s)
    def print_yellow(s):
        cprint(F_YELLOW, s)

    def print_boldred(s):
        cbprint(F_RED, s)
    def print_boldgreen(s):
        cbprint(F_GREEN, s)
    def print_boldblue(s):
        cbprint(F_BLUE, s)
    def print_boldmagenta(s):
        cbprint(F_MAGENTA, s)
    def print_boldcyan(s):
        cbprint(F_CYAN, s)

    def print_bold(s):
        bprint(s)

except:
    redstring = chr(27) + chr(91) + chr(51) + chr(49) + chr(109) + chr(0)
    normstring = chr(27) + chr(91) + chr(51) + chr(59) + chr(109) + chr(0)

    sys.stderr.write('\n#########################################################################\n')
    try:
        sys.stderr.write('%sGet pretty printing for this utility%s\n' \
                         % (redstring, normstring))
    except:
        sys.stderr.write('Get pretty printing for this utility\n')

    sys.stderr.write('Put \'/usr/groups/nprobe/jch1003/swig/utility/print_col.py on your PYPATH\'\n')
    sys.stderr.write('#########################################################################\n\n')

    def print_red(s):
        print(s)
    def print_green(s):
        print(s)
    def print_blue(s):
        print(s)
    def print_magenta(s):
        print(s)
    def print_cyan(s):
        print(s)
    def print_yellow(s):
        print(s)

    def print_boldred(s):
        print(s)
    def print_boldgreen(s):
        print(s)
    def print_boldblue(s):
        print(s)
    def print_boldmagenta(s):
        print(s)
    def print_boldcyan(s):
        print(s)

    def print_bold(s):
        print(s)
    
############################################################################## 
##############################################################################

def eprint(s):

    sys.stderr.write(s + '\n')
    
############################################################################## 
##############################################################################

def usage(nm, msg=None):

    if msg:
        eprint('%s: %s' % (nm, msg))
    
    eprint('%s: Find text or RE in Python file, giving context' % (nm))
    eprint('Usage:')
    eprint('\t%s <txt> [dir] - cwd if not dir given' % (nm))
    eprint('\t%s -r <regexp> [dir]' % (nm))
    eprint('Flags: \t[-h] this help\n\t[-d] recursively descend directory tree\n\t[-Dn] ditto to depth n\n\t[-s] use RE search for RE (default is match)\n\t[-f<file>] search for RE defined in file\n\t\t(One RE per line - multiple lines are treated as alternations)')
    eprint('\t[-l<file>|<list>] present numbered lines in context\n\t\t(lines as comma separated set of line numbers or ranges,\n\t\tor in file - one line number or range per line)')
    sys.exit(1)

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

def get_files(args,  recurse, ferr, matchfn, idre, classre, fnre, cmmtre, lines):

    for a in args:
        try:
            mode = os.stat(a)[ST_MODE]
        except OSError, s:
            ferr.append((a, 'stat', s))
            continue
        if S_ISDIR(mode) and recurse:
            args2 = []
            try:
                for f in os.listdir(a):
                    args2.append(join(a, f))
                get_files(args2,  recurse-1, ferr, matchfn, idre, classre, fnre, cmmtre, lines)
            except OSError, s:
                ferr.append((a, 'ls', s))
        elif S_ISREG(mode) and a[-3:] == '.py':
            #flist.append(a)
            do_file(a, ferr, matchfn, idre, classre, fnre, cmmtre, lines)

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

def do_file(fnm, ferr, matchfn, idre, classre, fnre, cmmtre, lines):
    
    module_printed = 0
    Class = None
    class_printed = 0
    Fn = None
    fn_printed = 0
    method = None
    method_printed = 0
    method_ind = 0
    submethod = None
    submethod_printed = 0
    submethod_ind = 0

    indok = 1
    badindstart = None

    try:
        f = open(fnm, 'r')
    except IOError, s:
        ferr.append((a, 'open', s))
        #print 'Couldn\'t open search file', s
        return
    lineno = 0
    module = os.path.basename(fnm)
    #print 'lines', lines
    for l in f.readlines():
        lineno += 1
        l = l[:-1]
        if (not l):
            continue
        #print l.count('\t'), 'tabs', len(l)-len(l.lstrip()), 'spaces'
        comment = 0
        if not cmmtre.match(l):
            # not a comment - look at syntax and indentation
            n = idre.match(l)
            if n:
                ind = len(n.group(1))
                #print 'indent', ind, lineno
            else:
                # print 'no ind match'
                ind = 0
            if ind == 0:
                # print 'cancelling Class, Fn, method, submethod line %d \'%s\'' % (lineno, l)
                Class = None
                Fn = None
                method = None
                submethod = None
            if submethod:
                if not submethod_ind:
                    submethod_ind = ind
                if  ind < submethod_ind:
                    submethod = None
                    submethod_ind = 0
            n = classre.match(l)
            if n:
                # print 'got class', Class
                Class = n.group(1)
                class_printed = 0
            n = fnre.match(l)
            if n:
                fnind = len(n.group(1))
                if fnind == 0:
                    Fn = n.group(2)
                    # print 'got fn', Fn, 'ind', fnind
                    fn_printed = 0
                    indok = 1
                elif (Class and method and fnind > method_ind) \
                         or (Fn and (ind > 0)):
                    submethod = n.group(2)
                    submethod_printed = 0
                    #submethod_ind = fnind
                    # print 'got submethod', submethod, 'ind', fnind
                    indok = 1
                elif Class:
                    Fn = None
                    method = n.group(2)
                    method_printed = 0
                    method_ind = fnind
                    submethod = None
                    # print 'got method', method, 'ind', fnind
                    indok = 1
                else:
                    #print 'XXX Don\'t understand indentation (%d) at %s line %d \'%s\' method ind = %d Class = %s' % (fnind, fnm, lineno, l, method_ind, Class)
                    indok = 0
                    badindstart = lineno
                    #break

                if indok == 0 and badindstart != None and badindstart != lineno:
                    sys.stderr.write('Didn\'t understand indentation lines %d - %d - ignored\n' % (badindstart, lineno))
                    badindstart = None

        else:
            comment = 1
           
        if matchfn((l, lineno)):
            if not module_printed:
                print_boldblue('MODULE %s' % (module))
                module_printed = 1
            if Class and not class_printed:
                print_boldgreen('CLASS %s' % (Class))
                class_printed = 1
            if method and not method_printed:
                print_boldred('METHOD %s' % (method))
                method_printed = 1
            if Fn and not fn_printed:
                print_boldmagenta('FUNCTION %s' % (Fn))
                fn_printed = 1
            if submethod and not submethod_printed:
                print_cyan('SUB-METHOD %s' % (submethod))
                submethod_printed = 1
            if not comment:
                print '%d: %s' % (lineno, l)
            else:
                print '%d:' % (lineno),
                print_yellow('%s' % (l))

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

def get_regex(fnm, snm):

    try:
        refile = open(fnm, 'r')
    except IOError, s:
        print 'Couldn\'t read regex file', s
        usage(snm)

    #re = [l[:-1] for l in refile.readlines() if len(l) > 1 and len(l[:-1].rstrip().lstrip()) > 0]

    # This one allows for finding lines containing only spaces
    re = [l[:-1] for l in refile.readlines() if len(l) > 1]
    #print re
    if len(re) == 0:
        eprint('No RE contained in file ' + fnm)
        sys.exit(1)

    if len(re) > 1:
        # get rid of duplicates
        rid = {}
        for r in re:
            if rid.has_key(r):
                continue
            rid[r] = 1
        re = rid.keys()
        re.sort()
        if len(re) > 1:
            # use alternation for multiple re's
            for r in re[1:]:
                re[0] += '|%s' % (r)

    #print 're is \'%s\'' % (re[0])

    return re[0]

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

def getranges(arg):

    def get_ranges(lines, arg):

        rangere = re.compile('(\d+)\s*-\s*(\d+)|(\d+)')
        ranges = []
        
        for lt in lines:
            error = 0
            l = lt[2]
            g =  rangere.match(l)
            if g:
                gg = g.groups()
                if gg[0] and gg[1]:
                    #print 'range', gg[0], gg[1]
                    ranges.append((int(gg[0]), int(gg[1])))
                elif gg[2]:
                    #print 'line', gg[2]
                    ranges.append((int(gg[2]), int(gg[2])))
                else:
                    error = 1
            else:
                error = 1
            if error:
                print 'whoops - malformed line specification: \'%s\' line/element %d \'%s\'' % (arg, lt[0], lt[1])
                return None

        return ranges

        

    def read_lines(fnm):

        try:
            f = open(fnm, 'r')
        except IOError, s:
            print s
            return None
        lines = []
        ln = 1
        for lhold in f.readlines():
            l = lhold.lstrip().rstrip()
            if not l:
                ln += 1
                continue
            lines.append((ln, lhold, l))
            ln += 1

        return lines

    def get_lines(arg):

        asplit = arg.split(',')
        lines = []
        part = 1
        for l in asplit:
            lines.append((part, l, l))
            part += 1
        return lines

    #
    # main fn starts here
    #
               
    try:
        mode = os.stat(arg)[ST_MODE]
        if S_ISREG(mode):
            txtranges = read_lines(arg)
    except OSError, s:
        if str(s).find('No such file or directory') >= 0:
            txtranges = get_lines(arg) 
        else:
            return None

    if txtranges:
        return get_ranges(txtranges, arg)



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

def has_line(lineno, lines):

    #print lineno, lines

    for ln in lines:
        if ln[0] <= lineno <= ln[1]:
            return 1

    return 0

##############################################################################
    
def main():

    recurse = None
    ferr = []
    scriptname = os.path.basename(argv[0])
    regex = None
    remeth = 'match'
    refile = None
    txt = None
    lines = []
    
    try:
        optlist, args = getopt.getopt(argv[1:], 'hrsf:D:dl:')

    except getopt.error, s:
        usage(scriptname, msg=str(s))

    for opt in optlist:
        if opt[0] == '-h':
            usage(scriptname)
        if opt[0] == '-r':
            try:
                regex = args[0]
                args = args[1:]
            except IndexError:
                usage(scriptname, 'No RE argument given')
        if opt[0] == '-s':
            remeth = 'search'
        if opt[0] == '-f':
            regex = get_regex(opt[1], scriptname)
        if opt[0] == '-d':
            recurse = 100000
        if opt[0] == '-l':
            lines = getranges(opt[1])
        if opt[0] == '-D':
            try:
                recurse = int(opt[1])
            except ValueError:
                usage(scriptname, 'Invalid recursion depth \'%s\'' % (opt[1]))
    
    if (not regex) and (not lines):
        if not len(args):
            usage(scriptname, 'No search expression given')
        txt = args[0]

        args = args[1:]
    
    if not args:
        args = [os.getcwd()]
        if not recurse:
            recurse = 1
    elif len(args) == 1 and isdir(args[0]) and not recurse:
        recurse = 1

    if regex:
        RE = re.compile(regex)
        #matchfn = getattr(RE, remeth)
        refn = getattr(RE, remeth)
        def mch(mfn):
            return lambda arg: mfn(arg[0])
        matchfn = mch(refn)
    elif lines:
        def mch(lines):
            return lambda arg: has_line(arg[1], lines)
        matchfn = mch(lines)
    else:
        def mch(txt):
            return lambda arg: (arg[0].find(txt) >= 0)
        matchfn = mch(txt)

    idre = re.compile('(\s*).*')
    classre = re.compile('class[ ]*(.*):.*')
    fnre = re.compile('([ ]*)def[ ]*(.*\(.*\)):.*')
    cmmtre = re.compile('\s*#+.*')

    get_files(args, recurse, ferr, matchfn, idre, classre, fnre, cmmtre, lines)

    if ferr:
        print
        ferr.sort()
        print_bold('Couldn\'t search the following files or directories:')
        for fe in ferr:
            print fe[0], ' - ', fe[2][1], '(' + fe[1] + ')'
            
             
##############################################################################


# Call main when run as script
if __name__ == '__main__':
        main()
