from gdb.printing import PrettyPrinter, register_pretty_printer
import gdb

class KthreadPrettyPrinter(object):

    def __init__(self, val):
        self.val = val

    def to_string(self):
        result = ""
        name = self.val["name"]
        state = self.val["state"]
        # Should get PC on stack for non-current thread otherwise will only get address in cpu_context_switch
        cpuState = self.val["cpuState"]
        cmdline = "info line *%s"%int(cpuState["eip"].cast(gdb.lookup_type("long").pointer()))
        line = gdb.execute(cmdline, to_string=True)
        #line = gdb.find_pc_line(int(cpuState["eip"].cast(gdb.lookup_type("long").pointer()).referenced_value()))
        currentThread = gdb.parse_and_eval("currentThread")
        if self.val == currentThread:
            result += "->"
        else:
            result += "  "
        result += "Addr: 0x%x, name: %s, state: %s, PC: %s" % (self.val, name, state, line)
        return result


class CustomPrettyPrinterLocator(PrettyPrinter):
    """Given a gdb.Value, search for a custom pretty printer"""

    def __init__(self):
        super(CustomPrettyPrinterLocator, self).__init__(
            "matos_pretty_printers", []
        )

    def __call__(self, val):
        """Return the custom formatter if the type can be handled"""

        typename = gdb.types.get_basic_type(val.type).tag
        if typename is None:
            typename = val.type.name

        if typename == "kthread":
            return KthreadPrettyPrinter(val)


class KthreadListDumpCmd(gdb.Command):
    """Prints the kthread list"""

    def __init__(self):
        super(KthreadListDumpCmd, self).__init__(
            "kthread_list_dump", gdb.COMMAND_USER
        )

    def _kthread_list_to_str(self, val):
        """Walk through the Kthread list.

        We will simply follow the 'next' pointers until we encounter the HEAD again
        """
        idx = 0
        head = val
        kthread_ptr = val
        result = ""
        while kthread_ptr != 0 and (idx == 0 or kthread_ptr != head):
            result += "\n%d: %s" % (idx, KthreadPrettyPrinter(kthread_ptr).to_string())
            kthread_ptr = kthread_ptr["next"]
            idx += 1
        result = ("Found a Linked List with %d kthread:" % idx) + result
        return result

    def complete(self, text, word):
        # We expect the argument passed to be a symbol so fallback to the
        # internal tab-completion handler for symbols
        return gdb.COMPLETE_SYMBOL

    def invoke(self, args, from_tty):
        # We can pass args here and use Python CLI utilities like argparse
        # to do argument parsing
        print("Args Passed: %s" % args)
        if args:
            kthread_ptr_val = gdb.parse_and_eval(args)
        else:
            kthread_ptr_val = gdb.parse_and_eval("currentThread")
        if str(kthread_ptr_val.type) != "struct kthread *":
            print("Expected pointer argument of type (struct kthread *)")
            return

        print(self._kthread_list_to_str(kthread_ptr_val))


class PhyMemDescListDumpCmd(gdb.Command):
    """Prints the memDesc list"""

    def __init__(self):
        super(PhyMemDescListDumpCmd, self).__init__(
            "phyMem_list_dump", gdb.COMMAND_USER)

    def _phy_memDesc_lit_to_str(self, val):
        result = ""
        idx = 0
        head = val
        memDesc_ptr = val
        prevAddr = None
        printBegin = True
        memAddrs = []
        while memDesc_ptr != 0 and (idx == 0 or memDesc_ptr != head):
            memAddrs.append(int(memDesc_ptr["phy_addr"]))
            memDesc_ptr = memDesc_ptr["next"]
            idx += 1

        memAddrs.sort()
        for currAddr in memAddrs:
            if printBegin:
                prevAddr = currAddr
                result += "[0x%x" % prevAddr
                printBegin = False
            if currAddr > prevAddr + 4096:
                result += "->0x%x[\n[0x%x" % ((prevAddr+4096), currAddr)
            prevAddr = currAddr
        if prevAddr:
            result += "->0x%x[\n" % (prevAddr+4096)
        result = ("Found a List with %d memDesc:\n" % idx) + result
        return result

    def complete(self, text, word):
        return gdb.COMPLETE_SYMBOL

    def invoke(self, args, from_tty):
        if args:
            desc = gdb.parse_and_eval(args)
        else:
            desc = gdb.parse_and_eval("phyFreePage")
        if str(desc.type) != "struct memDesc *":
            print("Expected pointer argument of type (struct memDesc *)")
            return

        print(self._phy_memDesc_lit_to_str(desc))


register_pretty_printer(None, CustomPrettyPrinterLocator(), replace=True)
KthreadListDumpCmd()
PhyMemDescListDumpCmd()