# -*- coding: utf-8 -*-

from collections import OrderedDict

from typhon.enum import makeEnum
from typhon.errors import userError
from typhon.nano.slots import NoAssignIR

"""
Static scope analysis, in several passes:
 * Discovering the static scope layout
 * Removing meta.context()
 * Removing meta.state()
 * Laying out specialized frames
 * Deslotification
"""

def scopeError(layout, message):
    return userError(u"Error in scope of %s: %s" % (layout.fqn, message))

# The scope of a name; where the name is defined relative to each name use.
SCOPE_OUTER, SCOPE_FRAME, SCOPE_LOCAL = makeEnum(u"scope",
    [u"outer", u"frame", u"local"])

# The severity of a name; how deeply-reified the binding and slot of a name
# are in the actual backing storage of a frame.
SEV_NOUN, SEV_SLOT, SEV_BINDING = makeEnum(u"severity",
    [u"noun", u"slot", u"binding"])

def layoutScopes(ast, outers, fqn, inRepl):
    """
    Perform scope analysis.
    """

    layoutPass = LayOutScopes(outers, fqn, inRepl)
    ast = layoutPass.visitExpr(ast)
    topLocalNames, localSize = layoutPass.top.collectTopLocals()
    return ast, layoutPass.top.outers, topLocalNames, localSize

LayoutIR = NoAssignIR.extend(
    "Layout", [],
    {
        "Expr": {
            "BindingExpr": [("name", "Noun"), ("layout", None)],
            "SlotExpr": [("name", "Noun"), ("layout", None)],
            "NounExpr": [("name", "Noun"), ("layout", None)],
            "TempNounExpr": [("name", "Noun"), ("layout", None)],
            "MetaContextExpr": [("layout", None)],
            "MetaStateExpr": [("layout", None)],
            "ObjectExpr": [("doc", None), ("patt", "Patt"),
                           ("auditors", "Expr*"), ("methods", "Method*"),
                           ("matchers", "Matcher*"), ("mast", None),
                           ("layout", None)],
            "-HideExpr": None,
        },
        "Patt": {
            "BindingPatt": [("name", "Noun"), ("layout", None)],
            "FinalPatt": [("name", "Noun"), ("guard", "Expr"),
                          ("layout", None)],
            "TempPatt": [("name", "Noun"), ("layout", None)],
            "VarPatt": [("name", "Noun"), ("guard", "Expr"), ("layout", None)],
        },
        "Matcher": {
            "MatcherExpr": [("patt", "Patt"), ("body", "Expr"),
                            ("layout", None)],
        },
        "Method": {
            "MethodExpr": [("doc", None), ("verb", None), ("patts", "Patt*"),
                           ("namedPatts", "NamedPatt*"), ("guard", "Expr"),
                           ("body", "Expr"), ("layout", None)],
        },
    }
)


def countLocalSize(lo, sizeSeen):
    sizeSeen = max(sizeSeen, lo.position + 1)
    for x in lo.children:
        sizeSeen = max(countLocalSize(x, sizeSeen), sizeSeen)
    return sizeSeen


class ScopeBase(object):
    position = -1

    def __init__(self, next, fqn):
        self.next = next
        self.children = []
        self.node = None
        self.fqn = fqn

    def addChild(self, child):
        if child is self:
            assert False, "BZZT WRONG"
        self.children.append(child)

    def findChild(self, name):
        """
        Search through children and find the given name.
        """

        for child in self.children:
            scope, idx, severity = child.find(name, False)
            if scope is not None:
                return scope, idx, severity
        raise scopeError(self, u"Impossibility in findChild")


class ScopeOuter(ScopeBase):
    def __init__(self, outers, fqn, inRepl):
        self.outers = OrderedDict()
        for i, outer in enumerate(outers):
            self.outers[outer] = i, SEV_SLOT
        self.children = []
        self.fqn = fqn
        self.inRepl = inRepl

    def collectTopLocals(self):
        # In an interactive context, we may want to keep locals defined at the
        # top level for future use.
        topLocals = [(u"", SEV_BINDING)] * 5
        scopeitems = self.children[:]
        numLocals = 0
        for sub in scopeitems:
            if isinstance(sub, ScopeItem):
                i = sub.position
                numLocals = max(numLocals, i + 1)
                while (i + 1) > len(topLocals):
                    topLocals.extend([(u"", SEV_BINDING)] * len(topLocals))
                topLocals[i] = sub.name, sub.severity
                scopeitems.extend(sub.children)

        return topLocals[:numLocals], countLocalSize(self, 0)

    def requireShadowable(self, name, toplevel):
        if name in self.outers and not (toplevel and self.inRepl):
            raise scopeError(self, u"Cannot redefine " + name)

    def deepen(self, name, severity):
        if name in self.outers:
            index, sev = self.outers[name]
            if sev.asInt < severity.asInt:
                self.outers[name] = index, severity

    def find(self, name, autoClose):
        try:
            index, severity = self.outers[name]
            return SCOPE_OUTER, index, severity
        except KeyError:
            return None, 0, None


class FrameTable(object):
    "Static frame layout information."

    __immutable_fields_ = 'dynamicGuards', 'frameInfo[*]', 'names'

    def __init__(self, frameInfo):
        self.frameInfo = frameInfo
        self.names = OrderedDict()
        self.dynamicGuards = OrderedDict()
        for i, (name, scope, _, severity) in enumerate(frameInfo):
            self.names[name] = i
            self.dynamicGuards[name] = i

    def positionOf(self, name):
        """
        The index of the position in the frame where the name would be placed,
        or -1 if no position is reserved.
        """

        return self.names.get(name, -1)

class ScopeFrame(ScopeBase):
    "Scope info associated with an object closure."

    def __init__(self, next, fqn):
        # Names closed over. The use of OrderedDict here forces this
        # dictionary to stay ordered according to insertion. Since we only
        # append to frameNames, iteration over this dict will be in the
        # correct order later. ~ C.
        self.frameNames = OrderedDict()
        # Names from outer scope used (not included in closure at runtime)
        self.outerNames = OrderedDict()
        return ScopeBase.__init__(self, next, fqn)

    def requireShadowable(self, name, toplevel):
        return self.next.requireShadowable(name, False)

    def deepen(self, name, severity):
        self.next.deepen(name, severity)

    def find(self, name, autoClose):
        scope, idx, severity = self.next.find(name, autoClose)
        if scope is SCOPE_OUTER:
            self.outerNames[name] = (idx, severity)
            return scope, idx, severity
        elif autoClose and name not in self.frameNames:
            self.frameNames[name] = (len(self.frameNames), scope, idx,
                                     severity)
            return SCOPE_FRAME, self.frameNames[name][0], severity
        elif name in self.frameNames:
            return SCOPE_FRAME, self.frameNames[name][0], severity
        else:
            return scope, idx, severity

    def computeFrameTable(self, nanopass):
        frameInfo = []
        for (n, (i, scope, idx, severity)) in self.frameNames.iteritems():
            if scope is None:
                nanopass.errorWithSpan(u"Name %s has no scope" % n, None)
            if severity is None:
                nanopass.errorWithSpan(u"Name %s has no severity" % n, None)
            frameInfo.append((n, scope, idx, severity))
        self.frameTable = FrameTable(frameInfo[:])


class ScopeBox(ScopeBase):
    "Scope info associated with a scope-introducing node."

    def __init__(self, next):
        ScopeBase.__init__(self, next, next.fqn)
        self.position = next.position

    def requireShadowable(self, name, toplevel):
        scope, idx, _ = self.find(name, False)
        if scope is SCOPE_OUTER:
            self.next.requireShadowable(name, False)

    def deepen(self, name, severity):
        self.next.deepen(name, severity)

    def find(self, name, autoClose):
        return self.next.find(name, autoClose)


class ScopeItem(ScopeBase):
    "A single name binding."
    def __init__(self, next, name, severity):
        self.name = name
        self.position = next.position + 1
        self.severity = severity
        return ScopeBase.__init__(self, next, next.fqn)

    def requireShadowable(self, name, toplevel):
        if self.name == name:
            raise scopeError(self, u"Cannot redefine " + name)
        self.next.requireShadowable(name, False)

    def deepen(self, name, severity):
        if self.name == name and self.severity.asInt < severity.asInt:
            self.severity = severity
        self.next.deepen(name, severity)

    def find(self, name, autoClose):
        if self.name == name:
            return SCOPE_LOCAL, self.position, self.severity
        return self.next.find(name, autoClose)


class LayOutScopes(NoAssignIR.makePassTo(LayoutIR)):
    """
    Set up scope boxes and collect variable definition sites.
    """
    def __init__(self, outers, fqn, inRepl=False):
        self.top = self.layout = ScopeOuter(outers, fqn, inRepl)

    def visitExprWithLayout(self, node, layout):
        origLayout = self.layout
        origLayout.addChild(layout)
        self.layout = layout
        result = self.visitExpr(node)
        layout.node = result
        self.layout = origLayout
        return result

    def visitExprNested(self, node):
        return self.visitExprWithLayout(node, ScopeBox(self.layout))

    def visitFinalPatt(self, name, guard, span):
        origLayout = self.layout
        self.layout.requireShadowable(name, True)
        result = self.dest.FinalPatt(name, self.visitExpr(guard), origLayout, span)
        # NB: If there's a guard, then we'll promote this to slot severity so
        # that auditors can recover the guard. ~ C.
        if isinstance(result.guard, self.dest.NullExpr):
            severity = SEV_NOUN
        else:
            severity = SEV_SLOT
        self.layout = ScopeItem(self.layout, name, severity)
        origLayout.addChild(self.layout)
        self.layout.node = result
        return result

    def visitTempPatt(self, name, span):
        origLayout = self.layout
        self.layout.requireShadowable(name, True)
        result = self.dest.TempPatt(name, origLayout, span)
        self.layout = ScopeItem(self.layout, name, SEV_NOUN)
        origLayout.addChild(self.layout)
        self.layout.node = result
        return result

    def visitVarPatt(self, name, guard, span):
        origLayout = self.layout
        self.layout.requireShadowable(name, True)
        result = self.dest.VarPatt(name, self.visitExpr(guard), origLayout, span)
        # Perhaps in the future we could do some sort of absorbing, but not
        # today. Nope.
        self.layout = ScopeItem(self.layout, name, SEV_SLOT)
        self.layout.node = result
        origLayout.addChild(self.layout)
        return result

    def visitBindingPatt(self, name, span):
        origLayout = self.layout
        self.layout.requireShadowable(name, True)
        result = self.dest.BindingPatt(name, origLayout, span)
        self.layout = ScopeItem(self.layout, name, SEV_BINDING)
        self.layout.node = result
        origLayout.addChild(self.layout)
        return result

    def visitHideExpr(self, body, span):
        return self.visitExprNested(body)

    def visitMethodExpr(self, doc, verb, patts, namedPatts, guard, body, span):
        origLayout = self.layout
        self.layout = ScopeBox(self.layout)
        origLayout.addChild(self.layout)
        result = self.dest.MethodExpr(
            doc, verb,
            [self.visitPatt(p) for p in patts],
            [self.visitNamedPatt(np) for np in namedPatts],
            self.visitExpr(guard), self.visitExpr(body), origLayout, span)
        self.layout.node = result
        self.layout = origLayout
        return result

    def visitMatcherExpr(self, patt, body, span):
        origLayout = self.layout
        self.layout = ScopeBox(self.layout)
        origLayout.addChild(self.layout)
        result = self.dest.MatcherExpr(self.visitPatt(patt),
                                       self.visitExpr(body), origLayout, span)
        self.layout.node = result
        self.layout = origLayout
        return result

    def visitObjectExpr(self, doc, patt, auditors, methods, matchers, mast, span):
        if isinstance(patt, self.src.IgnorePatt):
            objName = u'_'
        elif isinstance(patt, self.src.FinalPatt) or isinstance(
                patt, self.src.VarPatt):
            objName = patt.name
        else:
            objName = u'???'
        p = self.visitPatt(patt)
        origLayout = self.layout
        # Names defined in auditors exprs are visible inside the object but not
        # outside it, but aren't necessarily part of the frame.
        outerBox = ScopeBox(origLayout)
        origLayout.addChild(outerBox)
        auds = [self.visitExpr(a) for a in auditors]
        self.layout = ScopeFrame(outerBox, origLayout.fqn + u'$' + objName)
        # If we have auditors, then we need to make sure that the as-auditor
        # is added to the slot information, so we need to reify the slot.
        if auds and (isinstance(p, self.dest.FinalPatt) or
                     isinstance(p, self.dest.VarPatt)):
            self.layout.deepen(p.name, SEV_SLOT)
        outerBox.addChild(self.layout)
        result = self.dest.ObjectExpr(
            doc, p, auds,
            [self.visitMethod(m) for m in methods],
            [self.visitMatcher(m) for m in matchers],
            mast,
            # Everything else captures the layout previous to its node, but
            # here we store the ScopeFrame itself (since there's no other
            # good place to put it).
            self.layout, span)
        self.layout.node = result
        self.layout = origLayout
        return result

    def visitMetaContextExpr(self, span):
        return self.dest.MetaContextExpr(self.layout, span)

    def visitMetaStateExpr(self, span):
        return self.dest.MetaStateExpr(self.layout, span)

    def visitNounExpr(self, name, span):
        return self.dest.NounExpr(name, self.layout, span)

    def visitTempNounExpr(self, name, span):
        return self.dest.TempNounExpr(name, self.layout, span)

    def visitSlotExpr(self, name, span):
        self.layout.deepen(name, SEV_SLOT)
        return self.dest.SlotExpr(name, self.layout, span)

    def visitBindingExpr(self, name, span):
        self.layout.deepen(name, SEV_BINDING)
        return self.dest.BindingExpr(name, self.layout, span)

    def visitEscapeOnlyExpr(self, patt, body, span):
        origLayout = self.layout
        self.layout = ScopeBox(origLayout)
        origLayout.addChild(self.layout)
        p = self.visitPatt(patt)
        b = self.visitExpr(body)
        result = self.dest.EscapeOnlyExpr(p, b, span)
        self.layout.node = result
        self.layout = origLayout
        return result

    def visitEscapeExpr(self, ejPatt, ejBody, catchPatt, catchBody, span):
        origLayout = self.layout
        self.layout = layout1 = ScopeBox(origLayout)
        p = self.visitPatt(ejPatt)
        b = self.visitExpr(ejBody)
        self.layout = layout2 = ScopeBox(origLayout)
        origLayout.addChild(layout1)
        origLayout.addChild(layout2)
        cp = self.visitPatt(catchPatt)
        cb = self.visitExpr(catchBody)
        result = self.dest.EscapeExpr(p, b, cp, cb, span)
        layout1.node = result
        layout2.node = result
        self.layout = origLayout
        return result

    def visitFinallyExpr(self, body, atLast, span):
        return self.dest.FinallyExpr(
            self.visitExprNested(body),
            self.visitExprNested(atLast),
            span)

    def visitIfExpr(self, test, consq, alt, span):
        origLayout = self.layout
        self.layout = layout1 = ScopeBox(origLayout)
        origLayout.addChild(layout1)
        t = self.visitExpr(test)
        c = self.visitExpr(consq)
        self.layout = layout2 = ScopeBox(origLayout)
        origLayout.addChild(layout2)
        e = self.visitExpr(alt)
        result = self.dest.IfExpr(t, c, e, span)
        layout1.node = result
        layout2.node = result
        self.layout = origLayout
        return result

    def visitTryExpr(self, body, catchPatt, catchBody, span):
        b = self.visitExprNested(body)
        origLayout = self.layout
        self.layout = ScopeBox(origLayout)
        origLayout.addChild(self.layout)
        cp = self.visitPatt(catchPatt)
        cb = self.visitExpr(catchBody)
        result = self.dest.TryExpr(b, cp, cb, span)
        self.layout.node = result
        self.layout = origLayout
        return result


def bindNouns(ast):
    ast = ReifyMetaState().visitExpr(ast)
    ast = ReifyMetaContext().visitExpr(ast)
    ast = SpecializeNouns().visitExpr(ast)
    ast = ComputeFrameTables().visitExpr(ast)
    return ast


ReifyMetaStateIR = LayoutIR.extend(
    "ReifyMetaState", [], {
        "Expr": {
            "-MetaStateExpr": None,
        }
    }
)

class ReifyMetaState(LayoutIR.makePassTo(ReifyMetaStateIR)):

    def makeList(self, exprs, layout, span):
        _makeList = self.dest.NounExpr(u"_makeList", layout, span)
        return self.dest.CallExpr(_makeList, u"run", exprs, [], span)

    def makeNamePair(self, name, layout, span):
        k = self.dest.StrExpr(u"&&" + name, span)
        v = self.dest.BindingExpr(name, layout, span)
        pair = self.makeList([k, v], layout, span)
        return self.makeList([pair], layout, span)

    def visitMetaStateExpr(self, layout, span):
        s = layout
        while not isinstance(s, ScopeFrame):
            if isinstance(s, ScopeOuter):
                frame = {}
                break
            s = s.next
        else:
            frame = s.frameNames
        pairs = self.makeList([self.makeNamePair(name, layout, span)
                               for name in frame.keys()], layout, span)
        return self.dest.CallExpr(
            self.dest.NounExpr(u"_makeMap", layout, span), u"fromPairs",
            [pairs], [], span)


ReifyMetaContextIR = ReifyMetaStateIR.extend(
    "ReifyMetaContext", [], {
        "Expr": {
            "-MetaContextExpr": None,
        }
    }
)

class ReifyMetaContext(ReifyMetaStateIR.makePassTo(ReifyMetaContextIR)):

    def visitMetaContextExpr(self, layout, span):
        fqn = layout.fqn
        frame = ScopeFrame(layout, u'META')
        return self.dest.ObjectExpr(
            u"",
            self.dest.IgnorePatt(self.dest.NullExpr(span), span),
            [], [self.dest.MethodExpr(
                u"", u"getFQNPrefix", [], [], self.dest.NullExpr(span),
                self.dest.StrExpr(fqn + u'$', span), layout, span)],
            [], None, frame, span)


BoundNounsIR = ReifyMetaContextIR.extend(
    "BoundNouns", [], {
        "Expr": {
            "-NounExpr": None,
            "-TempNounExpr": None,
            "-SlotExpr": None,
            "-BindingExpr": None,
            "LocalExpr": [("name", "Noun"), ("index", None)],
            "FrameExpr": [("name", "Noun"), ("index", None)],
            "OuterExpr": [("name", "Noun"), ("index", None)],
        },
        "Patt": {
            "-TempPatt": None,
            "-FinalPatt": None,
            "-VarPatt": None,
            "NounPatt": [("name", "Noun"), ("guard", "Expr"),
                          ("index", None)],
            "FinalSlotPatt": [("name", "Noun"), ("guard", "Expr"),
                              ("index", None)],
            "VarSlotPatt": [("name", "Noun"), ("guard", "Expr"),
                            ("index", None)],
            "FinalBindingPatt": [("name", "Noun"), ("guard", "Expr"),
                                 ("index", None)],
            "VarBindingPatt": [("name", "Noun"), ("guard", "Expr"),
                               ("index", None)],
            "BindingPatt": [("name", "Noun"), ("index", None)],
        },
        "Method": {
            "MethodExpr": [("doc", None), ("verb", None), ("patts", "Patt*"),
                           ("namedPatts", "NamedPatt*"), ("guard", "Expr"),
                           ("body", "Expr"), ("localSize", None)],
        },
        "Matcher": {
            "MatcherExpr": [("patt", "Patt"), ("body", "Expr"),
                            ("localSize", None)],
        },
    }
)

class SpecializeNouns(ReifyMetaContextIR.makePassTo(BoundNounsIR)):

    def visitBindingPatt(self, name, layout, span):
        return self.dest.BindingPatt(name, layout.position + 1, span)

    def visitFinalPatt(self, name, guard, layout, span):
        _, _, severity = layout.findChild(name)
        guard = self.visitExpr(guard)
        if severity is SEV_NOUN:
            return self.dest.NounPatt(name, guard, layout.position + 1, span)
        elif severity is SEV_SLOT:
            return self.dest.FinalSlotPatt(name, guard, layout.position + 1, span)
        elif severity is SEV_BINDING:
            return self.dest.FinalBindingPatt(name, guard, layout.position + 1, span)
        else:
            assert False, "snape"

    def visitTempPatt(self, name, layout, span):
        _, _, severity = layout.findChild(name)
        guard = self.dest.NullExpr(span)
        if severity is SEV_NOUN:
            return self.dest.NounPatt(name, guard, layout.position + 1, span)
        elif severity is SEV_SLOT:
            return self.dest.FinalSlotPatt(name, guard, layout.position + 1, span)
        elif severity is SEV_BINDING:
            return self.dest.FinalBindingPatt(name, guard, layout.position + 1, span)
        else:
            assert False, "snape"

    def visitVarPatt(self, name, guard, layout, span):
        _, _, severity = layout.findChild(name)
        guard = self.visitExpr(guard)
        # NB: Not prepared to handle SEV_NOUN vars yet.
        if severity is SEV_SLOT:
            return self.dest.VarSlotPatt(name, guard, layout.position + 1, span)
        elif severity is SEV_BINDING:
            return self.dest.VarBindingPatt(name, guard, layout.position + 1, span)
        else:
            assert False, "snape"

    def makeStorage(self, name, index, scope, span):
        if scope is SCOPE_LOCAL:
            return self.dest.LocalExpr(name, index, span)
        elif scope is SCOPE_FRAME:
            return self.dest.FrameExpr(name, index, span)
        elif scope is SCOPE_OUTER:
            return self.dest.OuterExpr(name, index, span)
        else:
            assert False, "thesaurus"

    def bindingToSlot(self, noun, span):
        return self.dest.CallExpr(noun, u"get", [], [], span)

    def slotToNoun(self, noun, span):
        return self.dest.CallExpr(noun, u"get", [], [], span)

    def visitNounExpr(self, name, layout, span):
        # Automatically add to closures of frames.
        scope, idx, severity = layout.find(name, True)
        if scope is None:
            raise scopeError(layout, name + u" is not defined")
        storage = self.makeStorage(name, idx, scope, span)
        if severity is SEV_BINDING:
            return self.slotToNoun(self.bindingToSlot(storage, span), span)
        elif severity is SEV_SLOT:
            return self.slotToNoun(storage, span)
        else:
            return storage

    def visitTempNounExpr(self, name, layout, span):
        scope, idx, severity = layout.find(name, True)
        if scope is None:
            raise scopeError(layout, name + u" is not defined")
        storage = self.makeStorage(name, idx, scope, span)
        if severity is SEV_BINDING:
            return self.slotToNoun(self.bindingToSlot(storage, span), span)
        elif severity is SEV_SLOT:
            return self.slotToNoun(storage, span)
        else:
            return storage

    def visitSlotExpr(self, name, layout, span):
        scope, idx, severity = layout.find(name, True)
        if scope is None:
            raise scopeError(layout, name + u" is not defined")
        storage = self.makeStorage(name, idx, scope, span)
        if severity is SEV_BINDING:
            return self.bindingToSlot(storage, span)
        else:
            return storage

    def visitBindingExpr(self, name, layout, span):
        scope, idx, _ = layout.find(name, True)
        if scope is None:
            raise scopeError(layout, name + u" is not defined")
        return self.makeStorage(name, idx, scope, span)

    def visitMethodExpr(self, doc, verb, patts, namedPatts, guard, body,
                        layout, span):
        return self.dest.MethodExpr(
            doc, verb,
            [self.visitPatt(p) for p in patts],
            [self.visitNamedPatt(np) for np in namedPatts],
            self.visitExpr(guard),
            self.visitExpr(body),
            countLocalSize(layout, 0) + 2, span)

    def visitMatcherExpr(self, patt, body, layout, span):
        return self.dest.MatcherExpr(
            self.visitPatt(patt),
            self.visitExpr(body),
            countLocalSize(layout, 0) + 2, span)

class ComputeFrameTables(BoundNounsIR.selfPass()):
    """
    Visit every object's layout and compute their frame tables.
    """

    def visitObjectExpr(self, doc, patt, auditors, methods, matchers, mast,
                        layout, span):
        layout.computeFrameTable(self)
        return self.super.visitObjectExpr(self, doc, patt, auditors, methods,
                                          matchers, mast, layout, span)
