"""
Reference tracker for lltype data structures.
"""

import autopath, sys, os
import gc
from pypy.rpython.lltypesystem import lltype, llmemory
from pypy.rpython.memory.gcheader import header2obj
from pypy.translator.tool.reftracker import BaseRefTrackerPage, MARKER
from pypy.tool.uid import uid


class LLRefTrackerPage(BaseRefTrackerPage):

    def compute(self, objectlist, size_gc_header):
        self.size_gc_header = size_gc_header
        return BaseRefTrackerPage.compute(self, objectlist)

    def formatobject(self, o):
        lines = []
        for name, value in self.enum_content(o):
            if not isinstance(value, str):
                value = '0x%x' % uid(value)
            lines.append('%s = %s' % (name, value))
        s = '\n'.join(lines)
        t = shorttypename(lltype.typeOf(o))
        return t, s, ''

    def get_referrers(self, o):
        return []    # not implemented

    def get_referents(self, o):
        for name, value in self.enum_content(o):
            if not isinstance(value, str):
                yield value

    def edgelabel(self, o1, o2):
        slst = []
        for name, value in self.enum_content(o1):
            if value is o2:
                slst.append(name)
        return '/'.join(slst)

    def newpage(self, objectlist):
        return self.__class__(objectlist, self.size_gc_header)

    def normalize(self, o):
        if self.size_gc_header is not None:
            try:
                return header2obj[o]._obj
            except (KeyError, TypeError):
                pass
        return o

    def enum_content(self, o, name='', with_header=True):
        # XXX clean up
        T = lltype.typeOf(o)
        if (self.size_gc_header is not None and with_header
            and isinstance(T, lltype.ContainerType) and T._gckind == 'gc'):
            adr = llmemory.cast_ptr_to_adr(o._as_ptr())
            adr -= self.size_gc_header
            o = adr.get()._obj
            T = lltype.typeOf(o)
        if isinstance(T, lltype.Struct):
            try:
                gcobjptr = header2obj[o]
                fmt = '(%s)'
            except KeyError:
                gcobjptr = None
                fmt = '%s'
            for name in T._names:
                for name, value in self.enum_content(getattr(o, name), name,
                                                     with_header=False):
                    yield fmt % (name,), value
            if gcobjptr:
                GCT = lltype.typeOf(gcobjptr)
                if self.size_gc_header is not None:
                    for sub in self.enum_content(gcobjptr._obj,
                                                 with_header=False):
                        yield sub
                else:
                    # display as a link to avoid the same data showing up
                    # twice in the graph
                    yield 'header of', gcobjptr._obj
        elif isinstance(T, lltype.Array):
            for index, o1 in enumerate(o.items):
                for sub in self.enum_content(o1, str(index)):
                    yield sub
        elif isinstance(T, lltype.Ptr):
            if not o:
                yield name, 'null'
            else:
                yield name, self.normalize(lltype.normalizeptr(o)._obj)
        elif isinstance(T, lltype.OpaqueType) and hasattr(o, 'container'):
            T = lltype.typeOf(o.container)
            yield 'container', '<%s>' % (shorttypename(T),)
            for sub in self.enum_content(o.container, name, with_header=False):
                yield sub
        elif T == llmemory.Address:
            if not o:
                yield name, 'NULL'
            else:
                addrof = o.ref()
                T1 = lltype.typeOf(addrof)
                if (isinstance(T1, lltype.Ptr) and
                    isinstance(T1.TO, lltype.Struct) and
                    addrof._obj in header2obj):
                    yield name + ' @hdr', self.normalize(addrof._obj)
                else:
                    yield name + ' @', self.normalize(o.ptr._obj)
##                     if o.offset:
##                         yield '... offset', str(o.offset)
        else:
            yield name, str(o)

def shorttypename(T):
    return '%s %s' % (T.__class__.__name__, getattr(T, '__name__', ''))


def track(*ll_objects):
    """Invoke a dot+pygame object reference tracker."""
    lst = [MARKER]
    size_gc_header = None
    seen = {}
    for ll_object in ll_objects:
        if isinstance(ll_object, llmemory.GCHeaderOffset):
            size_gc_header = ll_object
            continue
        #if isinstance(lltype.typeOf(ll_object), lltype.Ptr):
        #    ptr = lltype.normalizeptr(ll_object)
        #    if ptr is not None:
        #        ll_object = ptr._obj
        #    else:
        #        ll_object = None
        if ll_object is not None and id(ll_object) not in seen:
            lst.append(ll_object)
            seen[id(ll_object)] = ll_object
    page = LLRefTrackerPage(lst, size_gc_header)
    # auto-expand one level, for now
    auto_expand = 1
    for i in range(auto_expand):
        page = page.content()
        for ll_object in lst[1:]:
            for name, value in page.enum_content(ll_object):
                if not isinstance(value, str) and id(value) not in seen:
                    lst.append(value)
                    seen[id(value)] = value
        page = page.newpage(lst)
    page.display()


if __name__ == '__main__':
    try:
        sys.path.remove(os.getcwd())
    except ValueError:
        pass
    T = lltype.GcArray(lltype.Signed)
    S = lltype.GcForwardReference()
    S.become(lltype.GcStruct('S', ('t', lltype.Ptr(T)),
                                  ('next', lltype.Ptr(S))))
    s = lltype.malloc(S)
    s.next = lltype.malloc(S)
    s.next.t = lltype.malloc(T, 5)
    s.next.t[1] = 123
    track(s)