import os
import os.path

from .NDSCard import NDSCard, NDSCardHeader, NitroFooter, NDSCardFNT
from .NDSCard import NDSCardBanner, NDSCardFAT
from .NitroFS import FileAllocationEntry, DirectoryTableEntry, EntryNameTableDirectoryEntry
from .NitroFS import EntryNameTableFileEntry, EntryNameTableEndOfDirectoryEntry
from .Utils import align

class Unpacker:
    """Callable for unpacking Nintendo DS ROMs. This is meant to be used with
    `shutil.unpack_archive`. This does not currently support unpacking ROMs that use overlays.
    """

    def __call__(self, *args, **kwargs):
        card = NDSCard(args[0])
        out_dir = args[1]
        out_path = lambda *s: os.path.join(out_dir, *s)

        if card.main_ovt or card.sub_ovt:
            raise NotImplementedError("ROMs with OVTs are not currently supported")

        with open(out_path('header.bin'), 'wb') as fp:
            fp.write(card.header.encode())
        with open(out_path('arm9.bin'), 'wb') as fp:
            fp.write(card.main_rom)
        with open(out_path('arm7.bin'), 'wb') as fp:
            fp.write(card.sub_rom)
        if hasattr(card, 'static_footer'):
            with open(out_path('static_footer.bin'), 'wb') as fp:
                fp.write(card.static_footer.encode())
        if hasattr(card, 'banner'):
            with open(out_path('banner.bin'), 'wb') as fp:
                fp.write(card.banner.encode())
        if hasattr(card, 'rsa_sig'):
            with open(out_path('rsa_sig.bin'), 'wb') as fp:
                fp.write(card.rsa_sig)

        # NitroFS
        dirnames = [None] * len(card.fnt.directory_table)
        dirnames[0] = 'data/'  # root directory name
        filenames = [None] * len(card.fat)
        Unpacker._get_files(
            dirnames, filenames, card.fnt.directory_table,
            card.fnt.entry_name_table, 0
        )
        for dirname_elems in (name.split('/') for name in dirnames):
            os.makedirs(out_path(*dirname_elems), exist_ok=True)
        for file_id in range(len(filenames)):
            with open(out_path(filenames[file_id]), 'wb') as fp:
                fp.write(card.file_data[file_id])


    @staticmethod
    def _get_files(
        dirnames: list[str],
        filenames: list[str],
        dir_table: NDSCardFNT.DirectoryTable,
        names_table: NDSCardFNT.EntryNameTable,
        dir_id: int
    ):
        """Get all directory names and filenames info into `dirnames` and `filenames`"""
        file_id = dir_table[dir_id].file_id
        for entry in names_table[dir_id]:
            if isinstance(entry, EntryNameTableEndOfDirectoryEntry):
                break
            if isinstance(entry, EntryNameTableDirectoryEntry):
                subdir_id = entry.directory_id ^ 0xf000
                parent_dir_id = dir_table[subdir_id].parent_id ^ 0xf000
                subdir_name = dirnames[parent_dir_id] + entry.entry_name + '/'
                dirnames[subdir_id] = subdir_name
                Unpacker._get_files(
                    dirnames, filenames, dir_table, names_table, subdir_id
                )
            else:  # EntryNameTableFileEntry
                path = dirnames[dir_id] + entry.entry_name
                filenames[file_id] = path
                file_id += 1


class Archiver:
    """Callable for repacking Nintendo DS ROMs that were unpacked by `NDSCard.Unpacker`. This is
    meant to be used with `shutil.make_archive`. This does not currently suppport repacking
    overlays.
    """

    def __call__(self, *args, **kwargs):
        out_filename = args[0]
        in_dir = args[1]
        in_path = lambda *s: os.path.join(in_dir, *s)
        card = NDSCard()
        files_start_offset = 0

        with open(in_path('header.bin'), 'rb') as fp:
            card.header = NDSCardHeader(fp.read())
        files_start_offset += card.header.header_size

        with open(in_path('arm9.bin'), 'rb') as fp:
            card.main_rom = fp.read()
        files_start_offset += len(card.main_rom)
        try:
            with open(in_path('static_footer.bin'), 'rb') as fp:
                card.static_footer = NitroFooter(fp.read())
            files_start_offset += len(card.static_footer)
        except FileNotFoundError:
            pass
        files_start_offset = align(files_start_offset, 0x200)

        with open(in_path('arm7.bin'), 'rb') as fp:
            card.sub_rom = fp.read()
        files_start_offset += len(card.sub_rom)
        files_start_offset = align(files_start_offset, 0x200)

        try:
            with open(in_path('banner.bin'), 'rb') as fp:
                card.banner = NDSCardBanner(fp.read())
        except FileNotFoundError:
            pass
        try:
            with open(in_path('rsa_sig.bin'), 'rb') as fp:
                card.rsa_sig = fp.read()
        except FileNotFoundError:
            pass

        # NitroFS
        dir_tree = list(os.walk(in_path('data/')))

        names_table = Archiver._make_names_table(dir_tree)
        num_files = 0
        for ent_entries in names_table:
            for ent_entry in ent_entries:
                files_start_offset += len(ent_entry)
                if isinstance(ent_entry, EntryNameTableFileEntry):
                    num_files += 1

        dir_table = Archiver._make_dir_table(names_table)
        length = len(dir_table) * len(DirectoryTableEntry())
        files_start_offset += length
        files_start_offset = align(files_start_offset, 0x200)

        # reserve space for FAT
        length = num_files * len(FileAllocationEntry())
        files_start_offset += length
        files_start_offset = align(files_start_offset, 0x200)

        try:
            files_start_offset += len(card.banner)  # banner
            files_start_offset = align(files_start_offset, 0x200)
        except NameError:
            pass

        card.fat = Archiver._make_fat(dir_tree, dir_table, files_start_offset)
        card.file_data = Archiver._get_files(dir_tree)

        card.fnt = NDSCardFNT()
        card.fnt.directory_table = dir_table
        card.fnt.entry_name_table = names_table

        with open(out_filename, 'wb') as fp:
            fp.write(card.encode())


    @staticmethod
    def _make_names_table(
        dir_tree: list[tuple[str, list[str], list[str]]]
    ) -> NDSCardFNT.EntryNameTable:
        out = NDSCardFNT.EntryNameTable()
        subdir_id = 0xf001
        for dir_info in dir_tree:
            dir_entry = []
            # process directories
            for dirname in dir_info[1]:
                dir_entry.append(
                    EntryNameTableDirectoryEntry(dirname, subdir_id)
                )
                subdir_id += 1
            # process files
            for filename in dir_info[2]:
                dir_entry.append(EntryNameTableFileEntry(filename))
            # done with this directory, add end entry
            dir_entry.append(EntryNameTableEndOfDirectoryEntry())
            out.append(dir_entry)
        # processed all directories
        return out


    @staticmethod
    def _make_dir_table(names_table: NDSCardFNT.EntryNameTable) -> NDSCardFNT.DirectoryTable:
        num_dirs = len(names_table)
        out = [DirectoryTableEntry() for _ in range(num_dirs)]
        out[0].parent_id = num_dirs  # root dir
        first_file_id = 0
        subdir_idx = 1

        for dir_idx in range(num_dirs):
            # store size of corresponding FNT names table
            # in `size` custom attribute
            out[dir_idx].size = 0
            out[dir_idx].file_id = first_file_id
            for entry in names_table[dir_idx]:
                out[dir_idx].size += len(entry)
                if isinstance(entry, EntryNameTableDirectoryEntry):
                    # assign parent_id for subdirs not yet processed
                    out[subdir_idx].parent_id = dir_idx | 0xf000
                    subdir_idx += 1
                elif isinstance(entry, EntryNameTableFileEntry):
                    first_file_id += 1

        offset = num_dirs * 8
        out[0].start = offset
        for dir_idx in range(1, num_dirs):
            offset += out[dir_idx-1].size
            out[dir_idx].start = offset

        # remove added `size` attributes
        for entry in out:
            del entry.size

        # convert to DirectoryTable
        out = NDSCardFNT.DirectoryTable(b''.join(
            entry.encode() for entry in out
        ))

        return out


    @staticmethod
    def _make_fat(
        dir_tree: list[tuple[str, list[str], list[str]]],
        dir_table: NDSCardFNT.DirectoryTable,
        files_start_offset: int
    ) -> NDSCardFAT:
        num_dirs = len(dir_table)
        out = NDSCardFAT()
        out.extend(
            FileAllocationEntry()
            for dir_info in dir_tree
            for file in dir_info[2]
        )
        offset = align(files_start_offset, 0x200)

        for dirtab_idx in range(num_dirs):
            file_id = dir_table[dirtab_idx].file_id
            for filename in dir_tree[dirtab_idx][2]:
                file_path = os.path.join(dir_tree[dirtab_idx][0], filename)
                length = os.stat(file_path).st_size
                out[file_id].start = offset
                out[file_id].end = offset + length
                offset = align(offset, 0x200)
                file_id += 1

        return out


    @staticmethod
    def _get_files(dir_tree: list[tuple[str, list[str], list[str]]]) -> list[bytes]:
        file_data = []
        for dir_info in dir_tree:
            for filename in dir_info[2]:
                with open(os.path.join(dir_info[0], filename), 'rb') as fp:
                    file_data.append(fp.read())
        return file_data