from zipfile import ZipFile, BadZipFile

from PyQt5 import QtGui as QG

from colours import TILE_PALETTES, palette_to_qt
from exceptions import ROMOpeningError


class NES_ROM(object):
    def __init__(self, filename):
        self.filename = filename

        # To track circular rotation of the PRG bytes, in case the
        # tiles aren't aligned to 16-byte boundaries and the user
        # moves it.
        self.initial_position = 0
        self.read_rom()

    def read_rom(self):
        if self.filename.lower().endswith('.nes'):
            with open(self.filename, 'rb') as rom:
                self.ines = rom.read()
        elif self.filename.lower().endswith('.zip'):
            try:
                with ZipFile(self.filename, 'r') as rom:
                    nesfiles = [f for f in rom.filelist if not f.is_dir()
                                and f.filename.lower().endswith(".nes")]
                    if len(nesfiles) != 1:
                        raise ROMOpeningError(
                            "Zipfile does not contain exactly one NES file"
                        )
                    self.ines = rom.open(nesfiles[0]).read()
            except BadZipFile as exc:
                raise ROMOpeningError(f"Problem opening zip file:\n{exc}")
        else:
            raise ROMOpeningError(
                f"Cannot deduce file type from file name:\n\n"
                f"{self.filename.split('/')[-1]}"
            )

        self.header = self.ines[0:16]
        body = self.ines[16:]
        PRG_size = self.header[4]*16384
        CHR_size = self.header[5]*8192

        self.PRG = body[0:PRG_size]
        if CHR_size == 0:
            # No chr, just read the whole ROM as tiles
            self.CHR = None
            self.tile_data = list(body)
        else:
            self.CHR = body[PRG_size:PRG_size+CHR_size]
            self.tile_data = list(self.CHR)

        self.tiles = [Tile(idx, self.tile_data)
                      for idx in range(0, len(self.tile_data)//16)]

    def update_tiles(self):
        for tile in self.tiles:
            tile.clear_caches()

    def rotate_tile_data(self, steps):
        """
        Rotates the tile data by given number of steps, either
        positively or negatively. The maximum rotation in either
        direction is 8 steps.
        """

        if steps > 0:
            steps = min(8 - self.initial_position, steps)
        elif steps < 0:
            steps = max(-8 - self.initial_position, steps)

        if steps == 0:
            return

        self.tile_data[:] = self.tile_data[steps:] + self.tile_data[:steps]
        self.initial_position += steps

    def byte_forward(self):
        self.rotate_tile_data(1)
        self.update_tiles()

    def byte_backward(self):
        self.rotate_tile_data(-1)
        self.update_tiles()


class Tile(object):
    """
    A tile is a view into a part of a NES ROM (either its PRG or its
    CHR). A tile just keeps an index into the ROM position by 16-byte
    offsets as well as individual palette information for this one
    view.

    """

    default_palette = TILE_PALETTES["Primary"]

    def __init__(self, index, tile_data):
        super().__init__()

        self.index = index
        self.tile_data = tile_data
        self.clear_caches()
        self.set_palette(self.default_palette.copy())

    def __repr__(self):
        return f"Tile(bytes={self.raw_tile}, palette={self.palette})"

    def __str__(self):
        num_to_ascii = {0: " ", 1: ".", 2: "o", 3: "#"}
        return "_"*10 + "\n" + "\n".join(
            "[" + "".join([num_to_ascii[val] for val in row]) + "]"
            for row in self.tile
        ) + "\n" + "-"*10

    def clear_caches(self):
        self._raw_tile = None
        self._tile = None
        self._pixmap = None

    def set_palette(self, new_palette):
        self.palette = new_palette.copy()
        self.update_pixmap()

    def update_pixmap(self):
        img_data = bytes(sum(self.tile, []))
        image = QG.QImage(img_data, 8, 8, QG.QImage.Format_Indexed8)
        image.setColorTable(palette_to_qt(self.palette))
        self._pixmap = QG.QPixmap(image)

    def get_pixmap(self):
        if self._pixmap is None:
            self.update_pixmap()
        return self._pixmap

    def set_pixmap(self, pixmap):
        # Not implemented
        pass

    pixmap = property(get_pixmap, set_pixmap, clear_caches, "Pixmap representation of tile")

    def get_raw_tile(self):
        if self._raw_tile is None:
            self._raw_tile = self.tile_data[16*self.index:16*(self.index + 1)]
        return self._raw_tile

    def set_raw_tile(self, raw_tile):
        self.clear_caches()
        self.tile_data[16*self.index:16*(self.index + 1)] = raw_tile

    raw_tile = property(get_raw_tile, set_raw_tile, clear_caches, "Raw tile bytes as found in ROM")

    def get_tile(self):
        if self._tile is None:
            self._tile = self.parse_tile(self.raw_tile)
        return self._tile

    def set_tile(self, tile):
        self.clear_caches()
        self.raw_tile = self.unparse_tile(tile)

    tile = property(get_tile, set_tile, clear_caches, "A tile parsed as an 8x8 array of ints")

    @classmethod
    def parse_tile(cls, tile):
        """Given a raw tile's bytes, convert it to an 8x8 array of integers."""
        lowplane, hiplane = tile[0:8], tile[8:16]
        rows = []
        for lowbyte, hibyte in zip(lowplane, hiplane):
            row = [
                ((lowbyte >> idx) & 1) + 2*((hibyte >> idx) & 1)
                for idx in range(7, -1, -1)
            ]
            rows.append(row)
        return rows

    @classmethod
    def unparse_tile(cls, rows):
        """Given an 8x8 array of integers, convert it to a raw tile's bytes."""
        lowplane = bytes([int(''.join(str(num & 1) for num in row), 2) for row in rows])
        hiplane = bytes([int(''.join(str((num & 2) >> 1) for num in row), 2) for row in rows])
        return lowplane + hiplane