# queries to do with module/class/function relationships
from __future__ import generators
from getTypeOf import getTypeOf, getTypeOfExpr
from bike.parsing.visitor import walk, walkSourceNodes
from bike.parsing.newstuff import getRegexMatchesInPythonPath, \
     generateModuleFilenamesInPythonPath, getSourceNode
from bike.query.common import MatchFinder, walkLinesContainingStrings, getScopeForLine
from bike import log
from bike.parsing.fastparserast import Module
import re

# first name is a child of scope
# (e.g. a method in a class, a nested function in a function, a class in a module)
def getChildByName(scope,name):
    matches = [c for c in scope.getChildNodes() if c.name == name]
    if matches != []:
        return matches[0]

def getAllClassesInHierarchy(root,classobj):
    print >>log.progress,"Getting root classes"
    rootClasses = getRootClassesOfHierarchy(classobj)
    relatedClasses = [] + rootClasses
    print >>log.progress,"Getting root classes - done"
    for rootClass in rootClasses:
        print >>log.progress,"Getting all derived classes for ",rootClass
        relatedClasses += getAllDerivedClasses(rootClass)
    return relatedClasses

def getRootClassesOfHierarchy(klass):
    if klass is None:  # i.e. dont have base class in our ast
        return None
    if klass.getBaseClassNames() == []:  # i.e. is a root class
        return[klass]
    else:
        rootclasses = []
        for base in klass.getBaseClassNames():
            baseclass = getTypeOf(klass,base)
            rootclass = getRootClassesOfHierarchy(baseclass)
            if rootclass is None:  # base class not in our ast
                rootclass = [klass]
            rootclasses+=rootclass
        return rootclasses


def getAllDerivedClasses(baseclass):
    state = DerivedClassSearchState()
    return getAllDerivedClasses_impl(baseclass,state)
    
def getAllDerivedClasses_impl(baseclass, state):
    #print "baseclass.name",baseclass.name
    for fname,potentualaliases in state.getCandidateFilesForName(baseclass.name):
        #print "fname,potentualaliases",fname,potentualaliases
        srcnode = getSourceNode(fname)
        walkLinesContainingStrings(srcnode.fastparseroot,
                                   ClassVisitor(baseclass,srcnode,state),
                                   [baseclass.name] + potentualaliases)    
    return state.subclasses


def getAllDerivedClasses_impl(baseclass, state):
    #print "getAllDerivedClasses_impl",baseclass
    for fname,classname in state.getPossibleDerivedClasses(baseclass.name):
        srcnode = getSourceNode(fname)
        walkLinesContainingStrings(srcnode.fastparseroot,
                                   ClassVisitor(baseclass,srcnode,state),
                                   ['class',classname])    
    return state.subclasses



classregex = re.compile("class\s+([^:( ]+)\s*\(([^)]+)\)")
wordregex = re.compile("\w+")
aliasregex = re.compile("from\s+[a-zA-Z\.]+\s+import\s+.+\w+\s+as\s+\w+")
asregex = re.compile("(\w+)\s+as\s+(\w+)")

class DerivedClassSearchState:
    def __init__(self):
        self.subclasses = []
        self.buildClassHierarchyMap()

    # maps classname -> (filename,classname)
    def buildClassHierarchyMap(self):
        print >>log.progress,"Building class hierarchy map of potentuals"
        map = {}

        for fname in generateModuleFilenamesInPythonPath():
            try:
                f = file(fname)
                src = f.read()
            finally:
                f.close()

            aliasMap = self._getAliasMap(src)
                    
            for match in classregex.finditer(src):
                for baseclassmatch in wordregex.finditer(match.group(2)):
                    baseclass = self._resolveAliases(aliasMap, baseclassmatch.group(0))
                    try:
                        map[baseclass].append((fname,match.group(1)))
                    except KeyError:
                        map[baseclass] = [(fname,match.group(1))]
        self.classhierarchymap = map
        print >>log.progress,"Building class hierarchy map of potentuals - Done"

    # returns list of (filename,classname)
    def getPossibleDerivedClasses(self,baseclassname):
        try:
            return self.classhierarchymap[baseclassname]
        except KeyError:
            return []

    def _resolveAliases(self, aliases, baseclass):
        try:
            baseclass = aliases[baseclass]
        except KeyError:
            pass
        return baseclass

    def _getAliasMap(self, src):
        # scan for aliases
        aliases = {}
        for match in aliasregex.finditer(src):
            for match in asregex.finditer(match.group(0)):
                aliases[match.group(2)] = match.group(1)
        return aliases

    # generates (filename,[potentual aliases])
    def getCandidateFilesForName(self,name):
        for filename, matchstrs in self.classesAndAliases.iteritems():
            aliases = []
            gotAMatch = False
            for matchstr in matchstrs:
                match = re.search(name+"\s+as\s+(\w+)",matchstr)
                if match is not None:
                    aliases.append(match.group(1))
                if matchstr.find(name) != -1:
                    gotAMatch = True
            if gotAMatch:
                yield filename,aliases



class ClassVisitor:
    def __init__(self,baseclass,srcnode,state):
        self.baseclass = baseclass
        self.state = state
        self.srcnode = srcnode
        
    def visitClass(self, node):
        classscope = getScopeForLine(self.srcnode, self.lineno)
        
        # create a list of possible aliases for the base class
        # (used to restrict type searches)
        baseclassname = self.baseclass.name
        targetBaseclassNames = [baseclassname] + \
                         scanImportsForAliases(classscope,baseclassname)

        classBaseClasses = classscope.getBaseClassNames()
        for basename in classBaseClasses:
            if basename.split('.')[-1] in targetBaseclassNames and \
                   getTypeOf(classscope,basename) == self.baseclass:
                self.state.subclasses.append(classscope)
                getAllDerivedClasses_impl(classscope,self.state)


# looks for import name as foo, returns [foo,..etc..]
def scanImportsForAliases(scope,name):
    aliases = []        
    while scope:
        for lineno in scope.getImportLines():
            line = scope.module.getSourceNode().getLines()[lineno-1]
            match = re.search(name+"\s+as\s+(\w+)",line)
            if match is not None:
                aliases.append(match.group(1))
        if not isinstance(scope,Module):
            scope = scope.getParent()
        else:
            scope = None
    return aliases
