import gdb
from kdumpfile import kdumpfile
from util import list_for_each_entry
import os
import os.path
import fnmatch

#arch = "i386:x86-64"
#
#setup = {
#    'i386:x86-64' : setup_thread_amd64,
#}

ulong_type = gdb.lookup_type('unsigned long')
try:
    rip = gdb.lookup_minimal_symbol("thread_return").value()
except:
    rip = gdb.lookup_minimal_symbol("schedule").value()

voidpp = gdb.lookup_type('void').pointer().pointer()
charp = gdb.lookup_type('char').pointer()

class PyKdumpCmd(gdb.Command):
    def __init__(self):
        gdb.Command.__init__(self, "pykdump", gdb.COMMAND_USER,gdb.COMPLETE_FILENAME)
    def invoke(self, argstr, from_tty):
        f = file(argstr, 'r')
        t = Target(f)

class PyKmodCmd(gdb.Command):
    def __init__(self):
        gdb.Command.__init__(self, "pykmod", gdb.COMMAND_USER,gdb.COMPLETE_FILENAME)
    def invoke(self, argstr, from_tty):
        mods = {}
        for root, dirnames, filenames in os.walk(argstr):
            for filename in fnmatch.filter(filenames, '*.ko.debug'):
                modname = os.path.basename(filename).rstrip(".ko.debug")
                mods [modname] = os.path.join(root, filename)

        list_head_type = gdb.lookup_type('struct list_head').pointer()
        modules_list = gdb.lookup_minimal_symbol('modules').value().cast(voidpp).cast(list_head_type).dereference()
        module_type = gdb.lookup_type('struct module')

        for mod in list_for_each_entry(modules_list, module_type, 'list'):
            modname = mod["name"].string()
            module_core = long(mod["module_core"].cast(voidpp))
            if not mods.has_key(modname):
                gdb.write("Debuginfo for %s not found (to be loaded at %lx)!\n" % (modname, module_core))
            else:
                gdb.execute("add-symbol-file %s 0x%lx" % (mods[modname], module_core))

            #gdb.execute("thread %d" % t.num)

class PydCmd(gdb.Command):
    def __init__(self):
        gdb.Command.__init__(self, "pyd", gdb.COMMAND_USER)
    def invoke(self, argstr, from_tty):
        pid = long(argstr)
        for t in gdb.selected_inferior().threads():
            if t.ptid[1] == pid:
                gdb.execute("thread %d" % t.num)
                return

class PyTestCmd(gdb.Command):
    def __init__(self):
        gdb.Command.__init__(self, "pytest", gdb.COMMAND_USER)
    def invoke(self, argstr, from_tty):

        init_task = gdb.lookup_global_symbol('init_task')
        task_list = init_task.value()['tasks']

        n = 0
        for taskg in list_for_each_entry(task_list, init_task.type, 'tasks'):
            for task in list_for_each_entry(taskg['thread_group'], init_task.type, 'thread_group'):
            #gdb.write("Thread %lx pid %d\n" % (long(task.address), task['pid']))
                n += 1

        gdb.write("tasks %d\n" % n)

            #thread = gdb.selected_inferior().new_thread((1, task['pid'], 0), taskinfo)

            #thread.name = task['comm'].string()


PyKdumpCmd()
PyKmodCmd()
PydCmd()
PyTestCmd()

class Arch:
    def setup_thread_active(self, thread, task):
        raise NotImplementedError
    def setup_thread_scheduled(self, thread, task):
        raise NotImplementedError

def setup_thread_active_amd64(thread, task):
    for (k, v) in task.regs.items():
        if k in ["gs_base", "orig_ax", "rflags", "fs_base"]:
            continue
        thread.registers[k].value = v

def setup_thread_scheduled_amd64(thread, task):
    rsp = task.task_struct['thread']['sp'].cast(ulong_type.pointer())
    rbp = rsp.dereference().cast(ulong_type.pointer())
    rbx = (rbp - 1).dereference()
    r12 = (rbp - 2).dereference()
    r13 = (rbp - 3).dereference()
    r14 = (rbp - 4).dereference()
    r15 = (rbp - 5).dereference()

    # The two pushes that don't have CFI info
#    rsp += 2

#    ex = in_exception_stack(rsp)
#    if ex:
#    print "EXCEPTION STACK: pid %d" % task['pid']

    thread.registers['rsp'].value = rsp
    thread.registers['rbp'].value = rbp
    thread.registers['rip'].value = rip
    thread.registers['rbx'].value = rbx
    thread.registers['r12'].value = r12
    thread.registers['r13'].value = r13
    thread.registers['r14'].value = r14
    thread.registers['r15'].value = r15
    thread.registers['cs'].value = 2*8
    thread.registers['ss'].value = 3*8

class Arch_x86_64(Arch):
    def __init__(self, specinfo):
        self.archident = "i386:x86-64"
        pass

    def setup_thread_active(self, thread, task):
        return setup_thread_active_amd64(thread, task)

    def setup_thread_scheduled(self, thread, task):
        return setup_thread_scheduled_amd64(thread, task)


class Arch_s390x(Arch):
    def __init__(self, specinfo):
        self.archident = "s390:64-bit"
        self.lowcore_loaded = False
        self.cpu_to_lowcore = {}
        self.pid_to_cpu = {}
        pass

    def ensure_lowcore(self):
        if self.lowcore_loaded: return 
        lcs = gdb.lookup_global_symbol("lowcore_ptr").value()
        for i in range(0, lcs.type.range()[1]+2):
            lc = lcs[i]
            if long(lc) == 0L: break
            cpu = long(lc["cpu_nr"])
            self.cpu_to_lowcore[cpu] = lc
            self.pid_to_cpu[long(lc["current_pid"])] = cpu

        self.lowcore_loaded = True

    def setup_thread_active(self, thread, task):
        # I'm afraid libkdumpfile doesn't have that now
        pass

    def setup_thread_scheduled(self, thread, task):
        self.ensure_lowcore()
        pid = long(task.task_struct['pid'])
        #print "pid=%d, %s" % (task.task_struct['pid'], self.pid_to_cpu.has_key(long(task.task_struct['pid'])))
        if self.pid_to_cpu.has_key(pid):
            lc = self.cpu_to_lowcore[self.pid_to_cpu[pid]]
            for i in range(0,16):
                thread.registers["r%d"%i].value = lc["gpregs_save_area"][i]

            thread.registers["pswa"].value = lc["psw_save_area"]["addr"]
            return
        ksp = task.task_struct['thread']['ksp'].cast(ulong_type.pointer())
        rip = (ksp+17).dereference().cast(ulong_type.pointer())
        thread.registers["pswa"].value = rip
        ksp = (ksp+18).dereference().cast(ulong_type.pointer())
        regmap = [("r%d"%i, i+3) for i in range(1,16)]
        for (regname, stackpos) in regmap:
            thread.registers[regname].value = (ksp + stackpos).dereference()

architectures = {
    "x86_64" : Arch_x86_64 ,
    "s390x"  : Arch_s390x }

def symbol_func(symname):
    ms = gdb.lookup_minimal_symbol(symname)
    if not ms:
        print ("Cannot lookup symbol %s" % symname)
        raise RuntimeError("Cannot lookup symbol %s" % symname)
    return long(ms.value())

class Target(gdb.Target):

    class ThreadInfo:
        def __init__(self, task_struct, active = False, cpu = -1, regs = None):
            self.task_struct = task_struct
            self.active = active
            self.cpu = cpu
            self.regs = regs

    def __init__(self, fil):
        if isinstance(fil, str):
            fil = file(fil)
            self.fil = fil
        print "kdump (%s)" % fil
        self.kdump = kdumpfile(fil)
        self.kdump.symbol_func = symbol_func
        self.kdump.vtop_init()
        self.setup_arch()
        super(Target, self).__init__()
        gdb.execute('set print thread-events 0')
        self.setup_percpu()
        self.setup_tasks()

    def setup_arch(self):
        arch = self.kdump.attr("arch")
        name = arch["name"]
        if not architectures.has_key(name):
            raise NotImplementedError("Architecture %s not supported yet" % name)

        self.arch = architectures[name](arch)
        gdb.write("Arch %s\n" % name)
        gdb.execute("set architecture %s" % self.arch.archident)

    def setup_tasks(self):
        init_task = gdb.lookup_global_symbol('init_task')
        task_list = init_task.value()['tasks']
        rqs = self.percpu_get("runqueues")
        rqscurrs = { long(x["curr"]) : k for (k, x) in rqs.items() }

        self.pid_to_task_struct = {}

        tasks = []
        for taskg in list_for_each_entry(task_list, init_task.type, 'tasks'):
            tasks.append(taskg)
            for task in list_for_each_entry(taskg['thread_group'], init_task.type, 'thread_group'):
                tasks.append(task)

        for task in tasks:
            if rqscurrs.has_key(long(task.address)):
                cpu = rqscurrs[long(task.address)]
                taskinfo = self.ThreadInfo(task\
                        , active=True\
                        , cpu=cpu\
                        , regs=self.kdump.attr("cpu." + str(cpu) + ".reg"))
            else:
                taskinfo = self.ThreadInfo(task)

            thread = gdb.selected_inferior().new_thread((1, task['pid'], 0), taskinfo)

            thread.name = task['comm'].string()

        gdb.selected_inferior().executing = False

    def setup_percpu(self):
        pco = long(gdb.lookup_global_symbol('__per_cpu_offset').value().address)
        self.percpu_offsets = {}
        try:
            for i in range(self.kdump.attr("cpu.number")):
                self.percpu_offsets[i] = long(gdb.Value(pco + i*voidpp.sizeof).cast(voidpp).dereference())
        except:
            pass

    def percpu_get(self, name, cpu=None):
        symbol = gdb.lookup_global_symbol(name)
        if not cpu == None:
            addr = long(symbol.value().address) + self.percpu_offsets[cpu]
            return gdb.Value(addr).cast(symbol.type.pointer()).dereference()
        else:
            return {c: gdb.Value(long(symbol.value().address) + adr)\
                .cast(symbol.type.pointer()).dereference() for (c, adr) in self.percpu_offsets.items()}

    def to_xfer_partial(self, obj, annex, readbuf, writebuf, offset, ln):
        ret = -1
        if obj == self.TARGET_OBJECT_MEMORY:
	    try:
		r = self.kdump.read (self.kdump.KDUMP_KVADDR, offset, ln)
		readbuf[:] = r
		ret = ln
	    except EOFError, e:
		raise gdb.TargetXferEof(str(e))
	    except IndexError, e:
		raise gdb.TargetXferUnavailable(str(e))
	    except:
		raise gdb.TargetXferUnavailable("nic")
	else:
	    raise IOError("Unknown obj type")
        return ret

    def to_thread_alive(self, ptid):
        return 1

    def to_pid_to_str(self, ptid):
        return "pid %d" % ptid[1]

    def to_fetch_registers(self, register):
        thread = gdb.selected_thread()
        if thread.info.active:
            self.arch.setup_thread_active(thread, thread.info)
        else:
            self.arch.setup_thread_scheduled(thread, thread.info)
        return True

    def to_prepare_to_store(self, thread):
        pass

    # We don't need to store anything; The regcache is already written.
    def to_store_registers(self, thread):
        pass

    def to_has_execution(self, ptid):
        return False