#
#  ktool | ktool
#  ktool.py
#
#  Outward facing API
#
#  Some of these functions are only one line long, but the point is to standardize an outward facing API that allows
#  me to refactor and change things internally without breaking others' scripts.
#
#  This file is part of ktool. ktool is free software that
#  is made available under the MIT license. Consult the
#  file "LICENSE" that is distributed together with this file
#  for the exact licensing terms.
#
#  Copyright (c) kat 2021.
#

from typing import Dict, Union, BinaryIO, List
from io import BytesIO

from ktool.dyld import Dyld, Image
from ktool.generator import TBDGenerator, FatMachOGenerator
try:
    from ktool.headers import HeaderGenerator, Header
except ModuleNotFoundError:
    # Maybe pygments wasn't installed and we're running in some weird context
    # So let whatever works, work
    Header = None
    HeaderGenerator = None
    pass
from ktool.macho import Slice, MachOFile, SlicedBackingFile
from ktool.objc import ObjCImage
from ktool.util import TapiYAMLWriter, ignore, log


def load_macho_file(fp: Union[SlicedBackingFile, BinaryIO, BytesIO], use_mmaped_io=True) -> MachOFile:
    """
    This function takes a bare file and loads it as a MachOFile.

    File should be opened with 'rb'

    :param fp: BinaryIO object
    :param use_mmaped_io: Should the MachOFile be loaded with a mmaped-io-backend? Leaving this enabled massively
                            improves load time and IO performance, only disable if your system doesn't support it
    :return:
    """
    if isinstance(fp, BytesIO):
        use_mmaped_io = False
    elif isinstance(fp, SlicedBackingFile):
        use_mmaped_io = False
        new_fp = BytesIO()
        new_fp.write(bytes(fp.file))
        new_fp.seek(0)
        fp = new_fp

    return MachOFile(fp, use_mmaped_io=use_mmaped_io)


def reload_image(image: Image) -> Image:
    """
    Reload an image (properly updates internal representations after patches)

    :param image:
    :return:
    """
    # This is going to be horribly slow. Dyld class needs refactored to have a better way to do this or ideally just
    #   not fuck things up and require a reload every time we make a patch.
    return load_image(image.slice)


def load_image(fp: Union[BinaryIO, MachOFile, Slice, BytesIO, SlicedBackingFile], slice_index=0, load_symtab=True, load_imports=True,
               load_exports=True, use_mmaped_io=True) -> Image:
    """
    Take a bare file, MachOFile, BytesIO, SlicedBackingFile, or Slice, and load MachO/dyld metadata about that item

    :param fp: a bare file, MachOFile, or Slice to load.
    :param slice_index: If a Slice is not being passed, and a file or MachOFile is a Fat MachO, which slice should be loaded?
    :param use_mmaped_io: If a bare file is being passed, load it with mmaped IO?
    :param load_symtab: Load the symbol table if one exists. This can be disabled for targeted loads, for speed.
    :param load_imports: Load imports if they exist. This can be disabled for targeted loads, for speed.
    :param load_exports: Load exports if they exist. This can be disabled for targeted loads, for speed.
    :return: Returns a loaded Image object
    :rtype: Image
    """
    if isinstance(fp, MachOFile):
        macho_file = fp
        macho_slice: Slice = macho_file.slices[slice_index]
    elif isinstance(fp, Slice):
        macho_slice = fp
    elif isinstance(fp, BytesIO):
        macho_file = load_macho_file(fp, use_mmaped_io=False)
        macho_slice: Slice = macho_file.slices[slice_index]
    elif isinstance(fp, SlicedBackingFile):
        macho_file = load_macho_file(fp, use_mmaped_io=False)
        macho_slice: Slice = macho_file.slices[slice_index]
    else:
        macho_file = load_macho_file(fp, use_mmaped_io=use_mmaped_io)
        macho_slice: Slice = macho_file.slices[slice_index]

    return Dyld.load(macho_slice, load_symtab=load_symtab, load_imports=load_imports, load_exports=load_exports)


def macho_verify(fp: Union[BinaryIO, MachOFile, Slice, Image]) -> None:
    """
    This function takes a variety of MachO-based objects, and loads them with malformation exceptions fully enabled.

    This can be used to verify patch code did not damage or improperly modify a MachO.

    :param fp: One of: BinaryIO, MachOFile, Slice, or Image, to load and verify
    :return:
    :raises: MalformedMachOException
    """
    should_ignore = ignore.MALFORMED

    log.info("Verifying MachO Integrity")
    ignore.MALFORMED = False

    if isinstance(fp, Image):
        load_image(fp.slice)

    elif isinstance(fp, MachOFile) or isinstance(fp, BinaryIO):
        if isinstance(fp, MachOFile):
            slices = fp.slices
        else:
            slices = load_macho_file(fp)

        for macho_slice in slices:
            load_image(macho_slice)

    else:
        load_image(fp)

    ignore.MALFORMED = should_ignore


def load_objc_metadata(image: Image) -> ObjCImage:
    return ObjCImage.from_image(image)


def generate_headers(objc_image: ObjCImage, sort_items=False) -> Dict[str, Header]:
    out = {}

    if sort_items:
        for objc_class in objc_image.classlist:
            objc_class.methods.sort(key=lambda h: h.signature)
            objc_class.properties.sort(key=lambda h: h.name)

        for objc_proto in objc_image.protolist:
            objc_proto.methods.sort(key=lambda h: h.signature)
            objc_proto.opt_methods.sort(key=lambda h: h.signature)

    for header_name, header in HeaderGenerator(objc_image).headers.items():
        out[header_name] = header

    return out


def generate_text_based_stub(image: Image, compatibility=True) -> str:
    generator = TBDGenerator(image, compatibility)
    return TapiYAMLWriter.write_out(generator.dict)


def macho_combine(slices: List[Slice]) -> BytesIO:
    fat_generator = FatMachOGenerator(slices)

    fat_file = BytesIO()
    fat_file.write(fat_generator.fat_head)

    for arch in fat_generator.fat_archs:
        fat_file.seek(arch.offset)
        fat_file.write(arch.slice.full_bytes_for_slice())

    fat_file.seek(0)
    return fat_file
