import os
from io import BytesIO, BufferedIOBase
from typing import overload, BinaryIO
from collections.abc import MutableSequence

from .Utils import uint_le, u8_le, u16_le, u32_le, align
from .CRC16 import crc16
from .NitroFS import FileAllocationEntry, EntryNameTableEndOfDirectoryEntry
from .NitroFS import EntryNameTableDirectoryEntry, DirectoryTableEntry
from .NitroFS import EntryNameTableEntry, EntryNameTableFileEntry

class NDSCard:
    def __init__(self, data: str|int|bytes|BinaryIO|None = None):
        """Read and parse a Nintendo DS ROM. If `data` is an open file object or file descriptor,
        this does not close it.
        """
        if data is None:
            self.header = NDSCardHeader()
            self.main_rom = b''
            self.sub_rom = b''
            self.fnt = NDSCardFNT()
            self.main_ovt = NDSCardOVT()
            self.sub_ovt = NDSCardOVT()
            self.fat = NDSCardFAT()
            self.file_data = []
            return

        if isinstance(data, bytes):
            data = BytesIO(data)
        elif isinstance(data, str):  # file path
            with open(data, 'rb') as fp:
                data = BytesIO(fp.read())
        elif isinstance(data, BufferedIOBase):  # open file object
            initial_fp = data
            initial_pos = data.tell()
            data = BytesIO(initial_fp.read())
            initial_fp.seek(initial_pos)
            del initial_fp, initial_pos
        elif isinstance(data, int):  # open file descriptor
            fp = os.fdopen(data, 'rb')
            initial_pos = fp.tell()
            data = BytesIO(fp.read())
            fp.seek(initial_pos)
            del initial_pos
        total_rom_size = data.seek(0, os.SEEK_END)
        data.seek(0)

        length = len(NDSCardHeader())
        self.header = NDSCardHeader(data.read(length))

        data.seek(self.header.main_rom_offset)
        self.main_rom = data.read(self.header.main_size)
        if uint_le(data.read(4)) == 0xDEC00621:  # Nitro footer
            data.seek(-4, os.SEEK_CUR)
            length = len(NitroFooter())
            self.static_footer = NitroFooter(data.read(length))

        data.seek(self.header.sub_rom_offset)
        self.sub_rom = data.read(self.header.sub_size)

        data.seek(self.header.fnt_offset)
        self.fnt = NDSCardFNT(data.read(self.header.fnt_size))

        data.seek(self.header.main_ovt_offset)
        self.main_ovt = NDSCardOVT(data.read(self.header.main_ovt_size))
        data.seek(self.header.sub_ovt_offset)
        self.sub_ovt = NDSCardOVT(data.read(self.header.sub_ovt_size))

        data.seek(self.header.fat_offset)
        self.fat = NDSCardFAT(data.read(self.header.fat_size))

        if self.header.banner_offset != 0:
            data.seek(self.header.banner_offset)
            length = len(NDSCardBanner())
            self.banner = NDSCardBanner(data.read(length))

        self.file_data = []
        for entry in self.fat:
            data.seek(entry.start)
            self.file_data.append(data.read(entry.end - entry.start))

        # RSA signature
        if self.header.rom_size+0x88 <= total_rom_size:
            data.seek(self.header.rom_size)
            rsa_sig = data.read(0x88)
            # check whether it's just padding
            for byte in rsa_sig:
                if byte not in (0xff, 0):
                    self.rsa_sig = rsa_sig
                    break


    def encode(self, multiboot: bool = False) -> bytes:
        out = BytesIO()

        # skip the header, and write it afterwards
        self.header.header_size = 0x4000
        out.write(bytes(0x4000))

        # ARM9 ROM
        self.header.main_rom_offset = out.tell()
        self.header.main_size = len(self.main_rom)
        out.write(self.main_rom)
        # static footer
        try:
            out.write(self.static_footer.encode())
        except AttributeError:
            pass

        # main OVT
        if self.main_ovt:
            # alignment
            pos = out.tell()
            out.write(bytes(align(pos, 0x200) - pos))
            self.header.main_ovt_offset = out.tell()
            length = len(NDSCardOVT.OverlayTableEntry())
            self.header.main_ovt_size = len(self.main_ovt) * length
            out.write(self.main_ovt.encode())
            for overlay in self.main_ovt:
                pos = out.tell()
                out.write(bytes(align(pos, 0x200) - pos))
                self.fat[overlay.file_id].start = out.tell()
                self.fat[overlay.file_id].end = out.tell() + len(self.file_data[overlay.file_id])
                out.write(self.file_data[overlay.file_id])

        else:
            self.header.main_ovt_offset = 0
            self.header.main_ovt_size = 0
        pos = out.tell()
        out.write(b'\xff' * (align(pos, 0x200) - pos))

        # ARM7 ROM
        self.header.sub_rom_offset = out.tell()
        self.header.sub_size = len(self.sub_rom)
        out.write(self.sub_rom)

        # sub OVT
        if self.sub_ovt:
            pos = out.tell()
            out.write(bytes(align(pos, 0x200) - pos))
            self.header.sub_ovt_offset = out.tell()
            length = len(NDSCardOVT.OverlayTableEntry())
            self.header.sub_ovt_size = len(self.sub_ovt) * length
            out.write(self.sub_ovt.encode())
            for overlay in self.sub_ovt:
                pos = out.tell()
                out.write(bytes(align(pos, 0x200) - pos))
                self.fat[overlay.file_id].start = out.tell()
                self.fat[overlay.file_id].end = out.tell() + len(self.file_data[overlay.file_id])
                out.write(self.file_data[overlay.file_id])

        else:
            self.header.sub_ovt_offset = 0
            self.header.sub_ovt_size = 0
        pos = out.tell()
        out.write(b'\xff' * (align(pos, 0x200) - pos))

        # FNT
        self.header.fnt_offset = out.tell()
        out.write(self.fnt.encode())
        self.header.fnt_size = out.tell() - self.header.fnt_offset
        pos = out.tell()
        out.write(b'\xff' * (align(pos, 0x200) - pos))

        # FAT
        self.header.fat_offset = out.tell()
        self.header.fat_size = len(self.fat) * 8
        # skip the FAT, and write it after writing file data
        out.write(bytes(self.header.fat_size))

        # banner
        if hasattr(self, 'banner'):
            pos = out.tell()
            out.write(b'\xff' * (align(pos, 0x200) - pos))
            self.header.banner_offset = out.tell()
            out.write(self.banner.encode())
        else:
            self.header.banner_offset = 0

        # files
        for i in range(
            self.header.main_ovt_size//32 + self.header.sub_ovt_size//32,
            len(self.file_data)
        ):
            pos = out.tell()
            out.write(b'\xff' * (align(pos, 0x200) - pos))
            self.fat[i].start = out.tell()
            self.fat[i].end = out.tell() + len(self.file_data[i])
            out.write(self.file_data[i])

        pos = out.tell()
        self.header.rom_size = pos
        out.write(bytes(align(pos, 4) - pos))

        # device size
        if not multiboot:
            capacity_size = self.header.rom_size
            capacity_size |= capacity_size >> 16
            capacity_size |= capacity_size >> 8
            capacity_size |= capacity_size >> 4
            capacity_size |= capacity_size >> 2
            capacity_size |= capacity_size >> 1
            capacity_size += 1
            if capacity_size <= 0x20000:
                capacity_size = 0x20000
            capacity = -18
            while capacity_size != 0:
                capacity_size >>= 1
                capacity += 1
            self.header.device_size = 0 if capacity < 0 else capacity
        else:
            self.header.device_size = 0x0b

        # RSA
        try:
            out.write(self.rsa_sig)
        except AttributeError:
            pass

        # write FAT
        out.seek(self.header.fat_offset)
        out.write(self.fat.encode())

        # write header
        out.seek(0)
        out.write(self.header.encode())

        out.seek(0)
        return out.read()

    __bytes__ = encode


    def __len__(self):
        return len(self.encode())


class NDSCardHeader:
    def __init__(self, data: bytes|None = None):
        if data is None:
            return
        self.game_title = data[:0xc].decode('ascii', errors='surrogateescape').rstrip('\0')
        self.game_code = data[0xc:0x10].decode('ascii', errors='surrogateescape').rstrip('\0')
        self.maker_code = data[0x10:0x12].decode('ascii', errors='surrogateescape').rstrip('\0')

        self.product_id = data[0x12:0x13]
        self.device_type = data[0x13:0x14]
        self.device_size = data[0x14]
        self.reserved_a = data[0x15:0x1d]
        self.region = data[0x1d:0x1e]
        self.game_version = data[0x1e]
        self.property = data[0x1f:0x20]

        self.main_rom_offset = uint_le(data[0x20:0x24])
        self.main_entry_address = uint_le(data[0x24:0x28])
        self.main_ram_address = uint_le(data[0x28:0x2c])
        self.main_size = uint_le(data[0x2c:0x30])

        self.sub_rom_offset = uint_le(data[0x30:0x34])
        self.sub_entry_address = uint_le(data[0x34:0x38])
        self.sub_ram_address = uint_le(data[0x38:0x3c])
        self.sub_size = uint_le(data[0x3c:0x40])

        self.fnt_offset = uint_le(data[0x40:0x44])
        self.fnt_size = uint_le(data[0x44:0x48])

        self.fat_offset = uint_le(data[0x48:0x4c])
        self.fat_size = uint_le(data[0x4c:0x50])

        self.main_ovt_offset = uint_le(data[0x50:0x54])
        self.main_ovt_size = uint_le(data[0x54:0x58])

        self.sub_ovt_offset = uint_le(data[0x58:0x5c])
        self.sub_ovt_size = uint_le(data[0x5c:0x60])

        self.cmd_settings = data[0x60:0x64]
        self.key1_cmd_settings = data[0x64:0x68]
        self.banner_offset = uint_le(data[0x68:0x6c])
        self.secure_crc = uint_le(data[0x6c:0x6e])
        self.secure_delay = uint_le(data[0x6e:0x70])

        self.main_autoload_done = uint_le(data[0x70:0x74])
        self.sub_autoload_done = uint_le(data[0x74:0x78])

        self.secure_disable = data[0x78:0x80]
        self.rom_size = uint_le(data[0x80:0x84])
        self.header_size = uint_le(data[0x84:0x88])
        self.reserved_b = data[0x88:0x94]
        self.nand_rom_offset = uint_le(data[0x94:0x96])
        self.nand_rw_offset = uint_le(data[0x96:0x98])
        self.reserved_c = data[0x98:0xc0]

        self.logo_data = data[0xc0:0x15c]
        self.logo_crc = uint_le(data[0x15c:0x15e])
        self.header_crc = uint_le(data[0x15e:0x160])


    def encode(self) -> bytes:
        self.logo_crc = crc16(self.logo_data)
        header_data = b''.join([
            self.game_title.ljust(0xc, '\0').encode('ascii', errors='surrogateescape'),
            self.game_code.ljust(4, '\0').encode('ascii', errors='surrogateescape'),
            self.maker_code.ljust(2, '\0').encode('ascii', errors='surrogateescape'),

            self.product_id,
            self.device_type,
            u8_le(self.device_size),
            self.reserved_a,
            self.region,
            u8_le(self.game_version),
            self.property,

            u32_le(self.main_rom_offset),
            u32_le(self.main_entry_address),
            u32_le(self.main_ram_address),
            u32_le(self.main_size),

            u32_le(self.sub_rom_offset),
            u32_le(self.sub_entry_address),
            u32_le(self.sub_ram_address),
            u32_le(self.sub_size),

            u32_le(self.fnt_offset),
            u32_le(self.fnt_size),

            u32_le(self.fat_offset),
            u32_le(self.fat_size),

            u32_le(self.main_ovt_offset),
            u32_le(self.main_ovt_size),

            u32_le(self.sub_ovt_offset),
            u32_le(self.sub_ovt_size),

            self.cmd_settings,
            self.key1_cmd_settings,
            u32_le(self.banner_offset),
            u16_le(self.secure_crc),
            u16_le(self.secure_delay),

            u32_le(self.main_autoload_done),
            u32_le(self.sub_autoload_done),

            self.secure_disable,
            u32_le(self.rom_size),
            u32_le(self.header_size),
            self.reserved_b,
            u16_le(self.nand_rom_offset),
            u16_le(self.nand_rw_offset),
            self.reserved_c,

            self.logo_data,
            u16_le(self.logo_crc)
        ])
        self.header_crc = crc16(header_data)

        return b''.join([
            header_data,
            u16_le(self.header_crc),
            bytes(0x1000 - 0x160)  # padding
        ])

    __bytes__ = encode


    def __len__(self):
        return 0x1000


class NitroFooter:
    def __init__(self, data: bytes|None = None):
        if data is None:
            self.nitro_code = 0
            self.module_params_offset = 0
            self.unknown = 0
            return
        self.nitro_code = uint_le(data[:4])
        self.module_params_offset = uint_le(data[4:8])
        self.unknown = uint_le(data[8:12])


    def encode(self) -> bytes:
        return b''.join([
            u32_le(self.nitro_code),
            u32_le(self.module_params_offset),
            u32_le(self.unknown)
        ])

    __bytes__ = encode


    def __len__(self):
        return 12


class NDSCardFNT:
    def __init__(self, data: bytes|None = None):
        if data is None:
            self.directory_table = NDSCardFNT.DirectoryTable()
            self.entry_name_table = NDSCardFNT.EntryNameTable()
            return
        num_dirs = uint_le(data[6:8])
        length = len(DirectoryTableEntry())
        self.directory_table = NDSCardFNT.DirectoryTable(data[:num_dirs*length])
        ent_offset = num_dirs * len(DirectoryTableEntry())
        self.entry_name_table = NDSCardFNT.EntryNameTable(data[ent_offset:], num_dirs)


    def encode(self) -> bytes:
        return b''.join([self.directory_table.encode(), self.entry_name_table.encode()])

    __bytes__ = encode


    def __len__(self):
        return len(self.encode())


    class DirectoryTable(MutableSequence):
        def __init__(self, data: bytes|None = None):
            self.entries: list[DirectoryTableEntry] = []
            if data is None:
                return
            length = len(DirectoryTableEntry())
            self.entries.append(DirectoryTableEntry(data[:length]))
            # the parent ID of the first directory table entry is
            # the total number of directories, including itself
            num_dirs = self.entries[0].parent_id
            for i in range(1, num_dirs):
                entry_data = data[i*length : (i*length)+length]
                self.entries.append(DirectoryTableEntry(entry_data))


        def encode(self) -> bytes:
            return b''.join(entry.encode() for entry in self.entries)

        __bytes__ = encode


        def __getitem__(self, index):
            return self.entries[index]


        def __setitem__(self, index, value):
            self.entries[index] = value


        def __delitem__(self, index):
            del self.entries[index]


        def __len__(self):
            return len(self.entries)


        def insert(self, index, value):
            self.entries.insert(index, value)


    class EntryNameTable(MutableSequence):
        @overload
        def __init__(self, data: bytes, num_dirs: int, /):
            ...

        @overload
        def __init__(self, data: None = None, /):
            ...

        def __init__(self, *args):
            self.entries: list[list[EntryNameTableEntry]] = []
            if not args or args[0] is None:
                return
            data: bytes = args[0]
            num_dirs: int = args[1]
            offset = 0
            isdir = 0x80
            end_entries = 0
            entries_in_dir = []
            while end_entries < num_dirs:
                entry_name_length = data[offset]
                if entry_name_length == 0:  # end entry
                    entries_in_dir.append(EntryNameTableEndOfDirectoryEntry())
                    self.entries.append(entries_in_dir)
                    entries_in_dir = []
                    end_entries += 1
                    offset += 1
                elif entry_name_length < isdir:  # file entry
                    entries_in_dir.append(
                        EntryNameTableFileEntry(data[offset : offset+entry_name_length+1])
                    )
                    offset += entry_name_length + 1
                else:  # directory entry
                    entries_in_dir.append(EntryNameTableDirectoryEntry(
                        data[offset : offset+(entry_name_length ^ isdir)+3]
                    ))
                    offset += (entry_name_length ^ isdir) + 3


        def encode(self) -> bytes:
            return b''.join(
                entry.encode()
                for entries in self.entries
                for entry in entries
            )

        __bytes__ = encode


        def __getitem__(self, index):
            return self.entries[index]

        def __setitem__(self, index, value):
            self.entries[index] = value

        def __delitem__(self, index):
            del self.entries[index]

        def __len__(self):
            return len(self.entries)

        def insert(self, index, value):
            self.entries.insert(index, value)


class NDSCardOVT(MutableSequence):
    def __init__(self, data: bytes|None = None):
        self.entries: list[NDSCardOVT.OverlayTableEntry] = []
        if data is None:
            return
        length = len(NDSCardOVT.OverlayTableEntry())
        for i in range(len(data) // length):
            self.entries.append(data[i*length : i*length+length])


    def encode(self) -> bytes:
        return b''.join(entry.encode() for entry in self.entries)

    __bytes__ = encode


    def __getitem__(self, index):
        return self.entries[index]

    def __setitem__(self, index, value):
        self.entries[index] = value

    def __delitem__(self, index):
        del self.entries[index]

    def __len__(self):
        return len(self.entries)

    def __bool__(self):
        return bool(self.entries)

    def insert(self, index, value):
        self.entries.insert(index, value)


    class OverlayTableEntry:
        def __init__(self, data: bytes|None = None):
            if data is None:
                self.id = 0
                self.ram_address = 0
                self.ram_size = 0
                self.bss_size = 0
                self.static_init_start = 0
                self.static_init_end = 0
                self.file_id = 0
                self.flags = NDSCardOVT.OverlayTableEntry.Flags()
                return
            self.id = uint_le(data[:0x4])
            self.ram_address = uint_le(data[0x4:0x8])
            self.ram_size = uint_le(data[0x8:0xc])
            self.bss_size = uint_le(data[0xc:0x10])
            self.static_init_start = uint_le(data[0x10:0x14])
            self.static_init_end = uint_le(data[0x14:0x18])
            self.file_id = uint_le(data[0x18:0x1c])
            self.flags = NDSCardOVT.OverlayTableEntry.Flags(data[0x1c:0x20])


        def encode(self) -> bytes:
            return b''.join([
                u32_le(self.id),
                u32_le(self.ram_address),
                u32_le(self.ram_size),
                u32_le(self.bss_size),
                u32_le(self.static_init_start),
                u32_le(self.static_init_end),
                u32_le(self.file_id),
                self.flags.encode()
            ])

        __bytes__ = encode


        def __len__(self):
            return 32


        class Flags:
            def __init__(self, data: bytes|None = None):
                if data is None:
                    self.compressed = 0
                    self.authentication_code = 0
                    return
                data = uint_le(data[:4])
                self.compressed = data & 0x00ffffff
                self.authentication_code = (data >> 24) & 0xff


            def encode(self) -> bytes:
                return u32_le(
                    0
                    | (self.compressed & 0x00ffffff)
                    | ((self.authentication_code & 0xff) << 24)
                )

            __bytes__ = encode


class NDSCardFAT(MutableSequence):
    def __init__(self, data: bytes|None = None):
        self.entries: list[FileAllocationEntry] = []
        if data is None:
            return
        length = len(FileAllocationEntry())
        self.entries.extend(
            FileAllocationEntry(data[i : i+length])
            for i in range(0, len(data), length)
        )


    def encode(self) -> bytes:
        return b''.join(entry.encode() for entry in self.entries)

    __bytes__ = encode


    def __getitem__(self, index):
        return self.entries[index]

    def __setitem__(self, index, value):
        self.entries[index] = value

    def __delitem__(self, index):
        del self.entries[index]

    def __len__(self):
        return len(self.entries)

    def __bool__(self):
        return bool(self.entries)

    def insert(self, index, value):
        self.entries.insert(index, value)


class NDSCardBanner:
    def __init__(self, data: bytes|None = None):
        if data is None:
            self.header = NDSCardBanner.BannerHeader()
            self.banner = NDSCardBanner.BannerV1()
            return
        header_length = len(NDSCardBanner.BannerHeader())
        banner_length = len(NDSCardBanner.BannerV1())
        self.header = NDSCardBanner.BannerHeader(data[:header_length])
        self.banner = NDSCardBanner.BannerV1(data[header_length : banner_length+header_length])


    def encode(self) -> bytes:
        self.header.crc16_v1 = crc16(self.banner.encode())
        return b''.join([self.header.encode(), self.banner.encode()])

    __bytes__ = encode


    def __len__(self):
        return 832


    class BannerHeader:
        def __init__(self, data: bytes|None = None):
            if data is None:
                self.version = 1
                self.crc16_v1 = 0
                self.reserved = bytes(28)
                return
            self.version = uint_le(data[0:2])
            self.crc16_v1 = uint_le(data[2:4])
            self.reserved = data[4:32]


        def encode(self) -> bytes:
            return b''.join([u16_le(self.version), u16_le(self.crc16_v1), self.reserved])

        __bytes__ = encode


        def __len__(self):
            return 32


    class BannerV1:
        image_size = 512
        palette_size = 32
        game_name_size = 256

        def __init__(self, data: bytes|None = None):
            if data is None:
                self.image = bytes(self.image_size)
                self.palette = bytes(self.palette_size)
                self.game_name = [''] * 6
                return

            offset = 0
            self.image = data[offset : offset+self.image_size]
            offset += self.image_size
            self.palette = data[offset : offset+self.palette_size]
            offset += self.palette_size
            self.game_name = []
            for _ in range(6):
                name = data[offset : offset+self.game_name_size].decode('utf_16_le').rstrip('\0')
                self.game_name.append(name)
                offset += self.game_name_size


        def encode(self) -> bytes:
            return b''.join([
                self.image,
                self.palette,
                b''.join(
                    name.ljust(self.game_name_size//2, '\0').encode('utf_16_le')
                    for name in self.game_name
                )
            ])

        __bytes__ = encode


        def __len__(self):
            return 800