from os.path import basename, dirname, splitext, join, isdir, isfile, getmtime
from os import mkdir, utime, walk
from io import BytesIO
import hashlib
import lxml.etree as etree
from time import time as unixtime

from kbinxml.kbinxml import KBinXML
from kbinxml.bytebuffer import ByteBuffer

from handlers import GenericFolder

SIGNATURE = 0x6CAD8F89
KBIN_OFFSET = 36

FILE_VERSION = 3

class IFS:
    def __init__(self, path):
        if isfile(path):
            self._load_ifs(path)
            self.is_file = True
        elif isdir(path):
            self._load_dir(path)
            self.is_file = False
        else:
            raise IOError('Input path does not exist')

    def _load_ifs(self, path):
        out = splitext(basename(path))[0] + '_ifs'
        self.default_out = join(dirname(path), out)
        self.ifs_out = path

        with open(path, 'rb') as f:
            self.file = f.read()
            b = ByteBuffer(self.file)

        signature = b.get_u32()
        if signature != SIGNATURE:
            raise IOError('Given file was not an IFS file!')
        self.file_version = b.get_u16()
        # next u16 is just NOT(version)
        assert b.get_u16() ^ self.file_version == 0xFFFF
        self.time = b.get_u32()
        self.tree_size = b.get_u32()
        self.header_end = b.get_u32()
        # 16 bytes for manifest md5, unchecked

        self.manifest = KBinXML(self.file[KBIN_OFFSET:])
        self._parse_manifest()

        assert self.tree_size == self._tree_size()

    def _load_dir(self, path):
        self.default_out = path
        self.ifs_out = basename(path).replace('_ifs', '.ifs')

        self.file_version = FILE_VERSION
        self.time = int(getmtime(path))
        self.tree_size = -1
        self.header_end = -1
        self.manifest = None

        os_tree = self._create_dir_tree(path)
        self.tree = GenericFolder.from_filesystem(self, os_tree)

    def _create_dir_tree(self, path):
        tree = self._create_dir_tree_recurse(walk(path))
        if 'ifs_manifest.xml' in tree['files']:
            tree['files'].remove('ifs_manifest.xml')

        return tree

    def _create_dir_tree_recurse(self, walker):
        tree = {}

        root, dirs, files = next(walker)
        tree['path'] = root
        tree['files'] = files
        tree['folders'] = []
        for dir in dirs:
            tree['folders'].append(self._create_dir_tree_recurse(walker))

        return tree

    def _parse_manifest(self):
        self.tree = GenericFolder.from_xml(self, self.manifest.xml_doc)

    def tostring(self):
        return self.tree.tostring()

    def extract_all(self, progress = True, recurse = True, path = None):
        self.out = path if path else self.default_out
        self._mkdir(self.out)
        if self.manifest:
            with open(join(self.out, 'ifs_manifest.xml'), 'wb') as f:
                f.write(self.manifest.to_text().encode('utf8'))
        self._extract_tree(self.tree, progress, recurse)

    def repack(self, progress = True, path = None):
        if path is None:
            path = self.ifs_out
        data_blob = BytesIO()

        self.manifest = KBinXML(etree.Element('imgfs'))
        manifest_info = etree.SubElement(self.manifest.xml_doc, '_info_')

        # the important bit
        self.tree.repack(self.manifest.xml_doc, data_blob, progress)
        data = data_blob.getvalue()

        data_md5 = etree.SubElement(manifest_info, 'md5')
        data_md5.attrib['__type'] = 'bin'
        data_md5.attrib['__size'] = '16'
        data_md5.text = hashlib.md5(data).hexdigest()

        data_size = etree.SubElement(manifest_info, 'size')
        data_size.attrib['__type'] = 'u32'
        data_size.text = str(len(data))

        manifest_bin = self.manifest.to_binary()
        self.header_end = 36 + len(manifest_bin)
        self.ifs_size = self.header_end + len(data)
        self.tree_size = self._tree_size()
        manifest_hash = hashlib.md5(manifest_bin).digest()

        head = ByteBuffer()
        head.append_u32(SIGNATURE)
        head.append_u16(self.file_version)
        head.append_u16(self.file_version ^ 0xFFFF)
        head.append_u32(int(unixtime()))
        head.append_u32(self.tree_size)
        head.append_u32(self.header_end)
        with open(path, 'wb') as ifs_file:
            ifs_file.write(head.data)
            ifs_file.write(manifest_hash)
            ifs_file.write(manifest_bin)
            ifs_file.write(data)

    # suspected to be the in-memory representation
    def _tree_size(self):
        BASE_SIZE = 856
        return BASE_SIZE + self._tree_size_recurse(self.tree)

    def _tree_size_recurse(self, tree, depth = 0):
        FILE = 64
        FOLDER = 56
        DEPTH_MULTIPLIER = 16

        size = len(tree.files) * FILE
        size += len(tree.folders) * (FOLDER - depth*DEPTH_MULTIPLIER)
        for name, folder in tree.folders.items():
            size += self._tree_size_recurse(folder, depth+1)
        return size

    def _extract_tree(self, tree, progress = True, recurse = True, dir = ''):
        outdir = join(self.out, dir)
        if progress:
            print(outdir)
        self._mkdir(outdir)
        for name, f in tree.files.items():
            out = join(outdir, f.name)
            if progress:
                print(out)

            data = f.load()
            self._save_with_time(out, data, f.time)
            if recurse and f.name.endswith('.ifs'):
                i = IFS(out)
                i.extract_all()

        for name, f in tree.folders.items():
            self._extract_tree(f, progress, recurse, join(dir, f.name))

        # fallback to file timestamp
        timestamp = tree.time if tree.time else self.time
        utime(outdir, (timestamp, timestamp))


    def _mkdir(self, dir):
        try:
            mkdir(dir)
        except FileExistsError:
            pass

    def load_file(self, start, size):
        start = self.header_end+start
        end = start + size
        assert start <= len(self.file) and end <= len(self.file)

        return self.file[start:end]

    def _save_with_time(self, filename, data, time):
        with open(filename, 'wb') as f:
            f.write(data)
        utime(filename, (time,time))

if __name__ == '__main__':
    import sys
    if len(sys.argv) < 2:
        print('ifstools filename.ifs OR folder_ifs')
        exit()
    i = IFS(sys.argv[1])
    if i.is_file:
        i.extract_all()
    else:
        i.repack()