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

"""
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.tadidx[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 = ' '

            if args.decompress and self.iscompressed(decdat):
                decdat = self.decompress(decdat)
            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)))

    def iscompressed(self, data):
        if len(data)<11:
            return
        size, flag = struct.unpack_from(">HH", data, 0)
        if size+5 != len(data):
            return
        if flag!=0x800:
            return
        if data[-3:] != b"\x00\x00\x02":
            return
        return True

    def decompress(self, data):
        C = zlib.decompressobj(-15)
        return C.decompress(data[8:-3])

def destruct_bank_definition(args, data):
    """
    decode the 'bank' / database definition
    """
    rd = ByteReader(data)

    version = rd.readbyte()
    print("bank version: %02x" % version)

    d = dict()
    while not rd.eof():
        keyname = rd.readname()
        if keyname in d:
            print("WARN: duplicate key: %s" % keyname)

        index_or_length = rd.readdword()
        if index_or_length >> 31:
            value = d[keyname] = rd.readbytes(index_or_length & 0x7FFFFFFF)
            if re.search(b'[^\x0d\x0a\x09\x20-\x7e\xc0-\xff]', value):
                print("%-20s - %s" % (keyname, toout(args, d[keyname])))
            else:
                print("%-20s - \"%s\"" % (keyname, strescape(d[keyname])))
        else:
            d[keyname] = index_or_length
            print("%-20s -> #%s" % (keyname, d[keyname]))
    return d

def decode_field(data):
    rd = ByteReader(data)
    typ = rd.readword()
    idx1 = rd.readdword()
    name = rd.readname()
    unk1 = rd.readdword()
    unk2 = rd.readbyte()
    if typ:
        idx2 = rd.readdword()
        unk3 = rd.readdword()
        unk4 = rd.readdword()
        remain = rd.readbytes()

        print("%d %2d/%2d %d,%d,%d,%d - '%s' -- %s" % (typ, idx1, idx2, unk1, unk2, unk3, unk4, name, tohex(remain)))
    else:
        print("%d %2d    %d,%d       - '%s'" % (typ, idx1, unk1, unk2, name))


def destruct_base_definition(args, data):
    """
    decode the 'base' / table definition
    """
    rd = ByteReader(data)

    version = rd.readbyte()
    print("base version: %02x" % version)
    unk123 = [rd.readword() for _ in range(3)]
    unk45 = [rd.readdword() for _ in range(2)]
    tablename = rd.readname()
    unkname = rd.readname()
    unk7 = rd.readdword()
    nrfields = rd.readdword()
    print("%d,%d,%d,%d,%d  %d,%d '%s'  '%s'" % (*unk123, *unk45, unk7, nrfields, tablename, unkname))
    fields = []
    for _ in range(nrfields):
        l = rd.readword()
        fielddef = rd.readbytes(l)
        if args.verbose:
            print("field: @%04x: %04x - %s" % (rd.o, l, tohex(fielddef)))
        fields.append(decode_field(fielddef))
    remaining = rd.readbytes()

    print("rem: %s" % tohex(remaining))
        


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:
                self.dumptabledefs(args)
                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 dumptabledefs(self, args):
        dbinfo = self.stru.readrec(1)
        dbdef = destruct_bank_definition(args, dbinfo)

        for i in range(100):
            idx = dbdef.get("Base%03d" % i)
            if idx:
                print("== Base%03d ==" % i)
                if type(idx)==int:
                    tbinfo = self.stru.readrec(idx)
                else:
                    # the table def is in the value.
                    tbinfo = struct.pack("<B", 4) + idx
                tbdef = destruct_base_definition(args, tbinfo)


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
        hexdump(args.offset, data, args)

    elif args.shift:
        # explicitly specified shift.
        args.shift = int(args.shift, 0)
        enc = kodecode(args.shift, data)
        hexdump(args.offset, enc, args)
    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.width:
        args.width = int(args.width, 0)
    else:
        args.width = 64 if args.ascdump else 16

    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)

    if args.type==1:
        destruct_bank_definition(args, data)
    elif args.type==2:
        destruct_base_definition(args, data)

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('--width', '-w', 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('--nodecompress', action='store_false', dest='decompress', default='true')
    cro.add_argument('dbdir', type=str)
    cro.set_defaults(handler=cro_dump)

    des = subparsers.add_parser('destruct', help='Stru dumper')
    des.add_argument('--verbose', '-v', action='store_true')
    des.add_argument('--ascdump', '-a', action='store_true')
    des.add_argument('--type', '-t', type=int, help='what type of record to destruct')
    des.set_defaults(handler=destruct)

    args = parser.parse_args()

    if args.handler:
        args.handler(args)


if __name__=='__main__':
    main()