import os.path
import io
import struct
from binascii import b2a_hex
from hexdump import hexdump, asasc, tohex, unhex
from koddecoder import kodecode
from readers import ByteReader

"""
python3 crodump.py crodump chechnya_proverki_ul_2012
python3 crodump.py kodump -s 6   -o 0x4cc9 -e 0x5d95 chechnya_proverki_ul_2012/CroStru.dat
"""

def toout(args, data):
    """ return either ascdump or hexdump """
    if args.ascdump:
        return asasc(data)
    else:
        return tohex(data)

def enumunreferenced(ranges, filesize):
    """ from a list of used ranges and the filesize, enumerate the list of unused ranges """
    o = 0
    for start, end, desc in sorted(ranges):
        if start > o:
            yield o, start-o
        o = end
    if o<filesize:
        yield o, filesize-o

class Datafile:
    """ Represent a single .dat with it's .tad index file """
    def __init__(self, dat, tad, need_decode=False):
        self.dat = dat
        self.tad = tad

        self.need_decode = need_decode

        self.readtad()

        self.dat.seek(0, io.SEEK_END)
        self.datsize = self.dat.tell()

    def readtad(self):
        self.tad.seek(0)
        hdrdata = self.tad.read(2*4)
        self.tadhdr = struct.unpack("<2L", hdrdata)
        indexdata = self.tad.read()
        self.tadidx = [ struct.unpack_from("<3L", indexdata, 12*_) for _ in range(len(indexdata)//12) ]

    def readdata(self, ofs, size):
        self.dat.seek(ofs)
        return self.dat.read(size)

    def readrec(self, idx):
        ofs, ln, chk = self.tabidx[idx-1]
        if ln==0xFFFFFFFF:
            # deleted record
            return

        flags = ln>>24

        ln &= 0xFFFFFFF
        dat = self.readdata(ofs, ln)

        if not dat:
            # empty record
            encdat = dat
        elif self.need_decode and not flags:
            extofs, extlen = struct.unpack("<LL", dat[:8])
            encdat = dat[8:]
            while len(encdat)<extlen:
                dat = self.readdata(extofs, 0x200)
                extofs, = struct.unpack("<L", dat[:4])
                encdat += dat[4:]
            encdat = encdat[:extlen]
        else:
            encdat = dat

        if not self.need_decode:
            return encdat
        else:
            return kodecode(idx, encdat)


    def dump(self, args):
        print("tadhdr: %08x %08x" % tuple(self.tadhdr))
        ranges = []  # keep track of used bytes in the .dat file.
        for i, (ofs, ln, chk) in enumerate(self.tadidx):
            if ln==0xFFFFFFFF:
                print("%5d: %08x %08x %08x" % (i+1, ofs, ln, chk))
                continue
            flags = ln>>24

            ln &= 0xFFFFFFF
            dat = self.readdata(ofs, ln)
            ranges.append((ofs, ofs+ln, "item #%d" % i))
            decflag = ' '
            infostr = ""
            tail = b''

            if not dat:
                # empty record
                encdat = dat
            elif self.need_decode and not flags:
                extofs, extlen = struct.unpack("<LL", dat[:8])
                infostr = "%08x;%08x" % (extofs, extlen)
                encdat = dat[8:]
                while len(encdat)<extlen:
                    dat = self.readdata(extofs, 0x200)
                    ranges.append((extofs, extofs+0x200, "item #%d ext" % i))
                    extofs, = struct.unpack("<L", dat[:4])
                    infostr += ";%08x" % (extofs)
                    encdat += dat[4:]
                tail = encdat[extlen:]
                encdat = encdat[:extlen]
                decflag = '+'
            else:
                encdat = dat
                decflag = '*'

            if self.need_decode:
                decdat = kodecode(i+1, encdat)
            else:
                decdat = encdat
                decflag = ' '

            print("%5d: %08x-%08x: (%02x:%08x) %s %s%s %s" % (i+1, ofs, ofs+ln, flags, chk, infostr, decflag, toout(args, decdat), tohex(tail)))

        if args.verbose:
            # output parts not referenced in the .tad file.
            for o, l in enumunreferenced(ranges, self.datsize):
                dat = self.readdata(o, l)
                print("%08x-%08x: %s" % (o, o+l, toout(args, dat)))


class Database:
    """ represent the entire database, consisting of stru, index and bank files """
    def __init__(self, dbdir):
        self.dbdir = dbdir

        self.stru = self.getfile("Stru", need_decode=True)
        self.index = self.getfile("Index")
        self.bank  = self.getfile("Bank")
        self.sys = self.getfile("Sys", need_decode=True)
        # BankTemp, Int


    def getfile(self, name, need_decode=False):
        try:
            datname = self.getname(name, "dat")
            tadname = self.getname(name, "tad")
            if datname and tadname:
                return Datafile(open(datname, "rb"), open(tadname, "rb"), need_decode)
        except IOError:
            return

    def getname(self, name, ext):
        """
        get a case-insensitive filename match for 'name.ext'.
        Returns None when no matching file was not found.
        """
        basename = "Cro%s.%s" % (name, ext)
        for fn in os.scandir(self.dbdir):
            if basename.lower() == fn.name.lower():
                return os.path.join(self.dbdir, fn.name)

    def dump(self, args):
        if self.stru:
            print("stru")
            self.stru.dump(args)
        if args.struonly:
            return
        if self.index:
            print("index")
            self.index.dump(args)
        if self.bank:
            print("bank")
            self.bank.dump(args)
        if self.sys:
            print("sys")
            self.sys.dump(args)


def incdata(data, s):
    """
    add 's' to each byte.
    This is useful for finding the correct shift from an incorrectly shifted chunk.
    """
    return b"".join(struct.pack("<B", (_+s)&0xFF) for _ in data)

def decode_kod(args, data):
    """
    various methods of hexdumping KOD decoded data.
    """
    if args.nokod:
        # plain hexdump, no KOD decode
        if args.ascdump:
            print(asasc(data))
        else:
            hexdump(args.offset, data)

    elif args.shift:
        # explicitly specified shift.
        args.shift = int(args.shift, 0)
        enc = kodecode(args.shift, data)
        hexdump(args.offset, enc)
    elif args.increment:
        # explicitly specified shift.
        for s in range(256):
            enc = incdata(data, s)
            print("%02x: %s" % (s, toout(args, enc)))
    else:
        # output with all possible 'shift' values.
        for s in range(256):
            enc = kodecode(s, data)
            print("%02x: %s" % (s, toout(args, enc)))


def kod_hexdump(args):
    """
    KOD decode a section of a data file
    """
    args.offset = int(args.offset, 0)
    if args.length:
        args.length = int(args.length, 0)
    elif args.endofs:
        args.endofs = int(args.endofs, 0)
        args.length = args.endofs - args.offset

    if args.filename:
        with open(args.filename, "rb") as fh:
            if args.length is None:
                fh.seek(0, io.SEEK_END)
                filesize = fh.tell()
                args.length = filesize-args.offset
            fh.seek(args.offset)
            data = fh.read(args.length)
            decode_kod(args, data)
    else:
        # no filename -> read from stdin.
        import sys
        data = sys.stdin.buffer.read()
        if args.unhex:
            data = unhex(data)
        decode_kod(args, data)


def cro_dump(args):
    """ handle 'crodump' subcommand """
    db = Database(args.dbdir)

    db.dump(args)


def destruct(args):
    """
    decode the index#1 structure information record
    Takes hex input from stdin.
    """
    import sys
    data = sys.stdin.buffer.read()
    data = unhex(data)

    d = dict()

    rd = ByteReader(data)

    o = 0
    startbyte = rd.readbyte()
    while not rd.eof():
        keylen = rd.readbyte()
        keyname = rd.readbytes(keylen).decode('ascii')
        if keyname in d:
            print("duplicate key: %s" % keyname)

        index_or_length = rd.readdword()
        if index_or_length >> 31:
            d[keyname] = rd.readbytes(index_or_length & 0x7FFFFFFF)
            print("%-20s - %s" % (keyname, toout(args, d[keyname])))
        else:
            d[keyname] = index_or_length
            print("%-20s -> %s" % (keyname, d[keyname]))

def main():
    import argparse
    parser = argparse.ArgumentParser(description='CRO hexdumper')
    subparsers = parser.add_subparsers()
    parser.set_defaults(handler=None)

    ko = subparsers.add_parser('kodump', help='KOD/hex dumper')
    ko.add_argument('--offset', '-o', type=str, default="0")
    ko.add_argument('--length', '-l', type=str)
    ko.add_argument('--endofs', '-e', type=str)
    ko.add_argument('--unhex', '-x', action='store_true', help="assume the input contains hex data")
    ko.add_argument('--shift', '-s', type=str, help="KOD decode with the specified shift")
    ko.add_argument('--increment', '-i', action='store_true', help="assume data is already KOD decoded, but with wrong shift -> dump alternatives.")
    ko.add_argument('--ascdump', '-a', action='store_true', help="CP1251 asc dump of the data")
    ko.add_argument('--nokod', '-n', action='store_true', help="don't KOD decode")
    ko.add_argument('filename', type=str, nargs='?', help="dump either stdin, or the specified file")
    ko.set_defaults(handler=kod_hexdump)

    cro = subparsers.add_parser('crodump', help='CROdumper')
    cro.add_argument('--verbose', '-v', action='store_true')
    cro.add_argument('--kodecode', '-k', action='store_true')
    cro.add_argument('--ascdump', '-a', action='store_true')
    cro.add_argument('--nokod', '-n', action='store_true')
    cro.add_argument('--struonly', action='store_true')
    cro.add_argument('dbdir', type=str)
    cro.set_defaults(handler=cro_dump)

    des = subparsers.add_parser('destruct', help='Stru dumper')
    des.add_argument('--ascdump', '-a', action='store_true')
    des.set_defaults(handler=destruct)

    args = parser.parse_args()

    if args.handler:
        args.handler(args)


if __name__=='__main__':
    main()