import glob
import importlib
import inspect
import os
import re
import sys
import itertools
import types
from filetypes.base import ParsingError, DummyAnalyzer
from filetypes.types import *
from malcat import UserTypeDefinition, Structure, FileType, Architecture
from dissect import cstruct


def parse_dynamic_types(data_dir, user_data_dir, reload_plugins=False):
    """
    return instances of Type python class (so not parsed)
    """
    if user_data_dir and user_data_dir != data_dir:
        toscan = (user_data_dir, data_dir)
    else:
        toscan = (data_dir,)
    seen = set()
    definitions = []
    lookup = {}
    for plugins_path in toscan:
        plugins_path = os.path.abspath(plugins_path)
        for fname in itertools.chain(
                glob.glob(os.path.join(plugins_path, "filetypes", "*.py")),
                glob.glob(os.path.join(plugins_path, "usertypes", "*.py")),
            ):
            module_name = os.path.relpath(fname, plugins_path)[:-3].replace(os.sep, ".")
            if module_name in sys.modules:
                mod = sys.modules[module_name]
                if reload_plugins:
                    for element in dir(mod):
                        if isinstance(element, types.ModuleType):
                            importlib.reload(element)
                    mod = importlib.reload(mod)
            else:
                mod = importlib.import_module(module_name)
            for elemname in dir(mod):
                elem = getattr(mod, elemname)
                try:
                    if issubclass(elem, Type) and elem != Type and not elemname in seen:
                        seen.add(elemname)
                        relpath = os.path.abspath(inspect.getfile(elem))
                        definition = UserTypeDefinition( ".".join([elem.__module__, elemname]), relpath)
                        lookup[definition.name] = elem
                        definitions.append(definition)
                except: pass
    return definitions, lookup


def fix_dissect_pointer_parsing(data):
    return re.sub(r"([a-zA-Z_][a-zA-Z0-9_]*)(?<!typedef)\s*(\*+)\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*;", r"\1 \2\3;", data)


def parse_static_types(data_dir, user_data_dir, architecture):
    """
    return instances of malcat.Structure (so already parsed)
    """
    if user_data_dir and user_data_dir != data_dir:
        toscan = (user_data_dir, data_dir)
    else:
        toscan = (data_dir,)
    is64 = architecture in (Architecture.X64,)
    seen = {
        "char": UInt8(),
        "uint8": UInt8(),
        "int8": Int8(),
        "wchar": UInt16(),
        "uint16": UInt16(),
        "int16": Int16(),
        "uint32": UInt32(),
        "int32": Int32(),
        "uint64": UInt64(),
        "int64": Int64(),
    }
    if is64:
        seen["void"] = UInt64()
        pointer_class = Va64
    else:
        seen["void"] = UInt32()
        pointer_class = Va32
    definitions = []
    lookup = {}
    for plugins_path in toscan:
        plugins_path = os.path.abspath(plugins_path)
        for fname in glob.glob(os.path.join(plugins_path, "usertypes", "*.h")):
            cparser = cstruct.cstruct()
            with open(fname, "r") as f:
                MalcatTokenParser(cparser).parse(f.read())
            for k, v in cparser.typedefs.items():
                if not isinstance(v, cstruct.types.structure.Structure) or k.startswith("_"):
                    continue
                v.name = k
                try:
                    field = resolve_type(v, seen, pointer_class)
                except:
                    import traceback
                    traceback.print_exc()
                    continue
                relpath = os.path.abspath(fname)
                typename = ".".join(["usertypes"] + [os.path.basename(fname).split(".")[0], k])
                definition = UserTypeDefinition(typename, relpath)
                lookup[typename] = Structure(Type.HEADER, k, field.native)
                definitions.append(definition)
    return sorted(definitions, key=lambda x: x.name), lookup


def resolve_type(dissect_struct, seen, pointer_class):
    """
    return binding.Field instance
    """
    if dissect_struct.name in seen:
        return seen[dissect_struct.name]
    struc = Struct(name=dissect_struct.name)
    for member in dissect_struct.lookup.values():
        field = None
        if isinstance(member.type, cstruct.types.base.Array):
            if type(member.type.count) != int:
                raise ValueError("Dynamic arrays not supported for {}.{}".format(dissect_struct.name, member.name))
            if member.type.type.name == "char":
                field = String(member.type.count, name=member.name)
            elif member.type.type.name == "wchar":
                field = StringUtf16le(member.type.count, name=member.name)
            else:
                subtype = resolve_type(member.type.type, seen, pointer_class)
                field = Array(member.type.count, subtype, name=member.name)
                field.native = malcat.ArrayField(subtype.native, member.type.count, "")
        elif isinstance(member.type, cstruct.types.pointer.Pointer):
            if member.type.type.name == "char":
                field = pointer_class(comment="pointer to {}".format(member.type.type.name), hint=String(0, zero_terminated=True))
            elif member.type.type.name == "wchar":
                field = pointer_class(comment="pointer to {}".format(member.type.type.name), hint=StringUtf16le(0, zero_terminated=True))
            elif member.type.type.name == "void":
                field = pointer_class()
            else:
                field = pointer_class(comment="pointer to {}".format(member.type.type.name))
        elif isinstance(member.type, cstruct.types.structure.Union):
            raise ValueError("Unions not supported for {}.{}".format(dissect_struct.name, member.name))
        elif isinstance(member.type, cstruct.types.structure.Structure):
            field = resolve_type(member.type, seen, pointer_class)
        else:
            field = seen.get(member.type.name, None)
        if field is None:
            raise KeyError("Cannot translate field type '{}' for {}.{}".format(member.type, dissect_struct.name, member.name))
        struc.native.add(member.name.strip(), field.native)
    seen[dissect_struct.name] = struc
    return struc


def instanciate_dynamic(offset, struc_class, analyzer, thefile):
    struc = None
    if analyzer is None:
        analyzer = DummyAnalyzer()
        analyzer.run(thefile)
    old_parsed = analyzer._parsed

    def ositer(e):
        yield e()

    try:
        analyzer._parsed = FileType(thefile, analyzer.__class__.name, analyzer.__class__.category)
        analyzer._offset = offset
        analyzer._stop = thefile.size
        analyzer._remaining = thefile.size - offset
        iterator = analyzer._iterate(ositer(struc_class))
        element = iterator.send(None)
        struc = Structure(Type.HEADER, "{} instance at #{:x}".format(element.__class__.__name__, offset), element.native)
        analyzer._parsed.add_structure(offset, struc)
        fa = analyzer._parsed.at(0)
        iterator.send(fa.value)
    except StopIteration:
        pass    
    finally:
        analyzer._parsed = old_parsed
    return struc





# includes some fixes for dissect.cstruct
class MalcatTokenParser(cstruct.parser.TokenParser):

    def _typedef(self, tokens):
        tokens.consume()
        type_ = None

        if tokens.next == self.TOK.IDENTIFIER:
            ident = tokens.consume()
            type_ = self.cstruct.resolve(ident.value)
        elif tokens.next == self.TOK.STRUCT:
            # The register thing is a bit dirty
            # Basically consumes all NAME tokens and
            # registers the struct
            type_ = self._struct(tokens, register=True)


        names = self._names(tokens)
        for name in names:
            while name.startswith("*"):
                name = name[1:]
                type_ = cstruct.types.pointer.Pointer(self.cstruct, type_)
            self.cstruct.addtype(name, type_)

    @staticmethod
    def _tokencollection():
        TOK = cstruct.parser.TokenCollection()
        TOK.add(r'#\[(?P<values>[^\]]+)\](?=\s*)', 'CONFIG_FLAG')
        TOK.add(r'#define\s+(?P<name>[^\s]+)\s+(?P<value>[^\r\n]+)\s*', 'DEFINE')
        TOK.add(r'typedef(?=\s)', 'TYPEDEF')
        TOK.add(r'(?:struct|union)(?=\s|{)', 'STRUCT')
        TOK.add(r'(?P<enumtype>enum|flag)\s+(?P<name>[^\s:{]+)\s*(:\s'
                r'*(?P<type>[^\s]+)\s*)?\{(?P<values>[^}]+)\}\s*(?=;)', 'ENUM')
        TOK.add(r'(?<=})\s*(?P<defs>(?:[a-zA-Z0-9_]+\s*,\s*)+[a-zA-Z0-9_]+)\s*(?=;)', 'DEFS')
        TOK.add(r'(?P<name>\*?\s*[a-zA-Z0-9_]+)(?:\s*:\s*(?P<bits>\d+))?(?:\[(?P<count>[^;\n]*)\])?\s*(?=;)', 'NAME')
        TOK.add(r'[a-zA-Z_][a-zA-Z0-9_]*', 'IDENTIFIER')
        TOK.add(r'[{}]', 'BLOCK')
        TOK.add(r'\$(?P<name>[^\s]+) = (?P<value>{[^}]+})\w*[\r\n]+', 'LOOKUP')
        TOK.add(r';', 'EOL')
        TOK.add(r'\s+', None)
        TOK.add(r'.', None)

        return TOK

    def _names(self, tokens):
        names = []
        while True:
            if tokens.next == self.TOK.EOL:
                tokens.eol()
                break

            if tokens.next not in (self.TOK.NAME, self.TOK.DEFS):
                break

            ntoken = tokens.consume()
            if ntoken == self.TOK.NAME:
                names.append(ntoken.value.replace(" ", ""))
            elif ntoken == self.TOK.DEFS:
                for name in ntoken.value.strip().split(','):
                    names.append(name.replace(" ", ""))

        return names


if __name__ == "__main__":
    #defs, lookup = parse_static_types("../data", "", None)
    print(lookup)
