#!/usr/bin/python3
import sys

from PyQt5 import QtCore as QC, QtGui as QG, QtWidgets as QW

from colors import (
    NES_PALETTE, TILE_PALETTES, is_dark, palette_to_qt,
    widget_icon_path
)


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

    def read_rom(self):
        with open(self.filename, 'rb') as f:
            self.ines = f.read()

        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
            tile_data = body
            self.CHR = None
        else:
            self.CHR = body[PRG_size:PRG_size+CHR_size]
            tile_data = self.CHR

        self.tiles = [Tile(tile_data[i:i+16]) for i in range(0, len(tile_data), 16)]


class Tile(object):

    default_palette = TILE_PALETTES["Primary"]

    def __init__(self, raw_tile):
        super().__init__()

        self.raw_tile = raw_tile
        self.set_palette(self.default_palette.copy())

    def __repr__(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 set_palette(self, new_palette):
        self.palette = new_palette
        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_tile(self):
        if getattr(self, "_tile", None) is None:
            self._tile = self.parse_tile(self.raw_tile)
        return self._tile

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

        self.update_pixmap()

    def del_tile(self):
        del self._tile
        del self.raw_tile

    tile = property(get_tile, set_tile, del_tile, "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


class GridHolder(object):
    def __init__(self, things, numrows, numcols, fillvalue=None):
        self.numrows = numrows
        self.numcols = numcols
        self.fillvalue = fillvalue
        if things is None:
            self.things = []
            for i in range(self.numrows):
                self.things.extend([fillvalue]*self.numcols)
        else:
            self.things = things

    def __getitem__(self, idx):
        i, j = idx
        try:
            return self.things[i*self.numcols + j]
        except IndexError:
            return self.fillvalue

    def __setitem__(self, idx, value):
        i, j = idx
        self.things[i*self.numcols + j] = value

    def __repr__(self):
        return f"GridHolder(numrows={self.numrows}, numcols={self.numcols}, fillvalue={self.fillvalue})"

    def __iter__(self):
        for thing in self.things:
            yield thing

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

    def resize(self, numrows, numcols):
        self.numrows = numrows
        self.numcols = numcols


class TileGrid(QW.QWidget):

    DEFAULT_SCALE_FACTOR = 5

    def __init__(self, numrows, numcols, scalefactor, spacing):
        super().__init__()

        self.numcols = numcols
        self.numrows = numrows
        self.scalefactor = scalefactor
        self.spacing = spacing

        self.setStyleSheet(self.CSS)

    def paintEvent(self, event):
        painter = QG.QPainter(self)

        # I don't really know what this block of code means, but it's
        # what I have to do in order to get this widget to obey CSS
        # styles.
        opt = QW.QStyleOption()
        opt.initFrom(self)
        style = self.style()
        style.drawPrimitive(QW.QStyle.PE_Widget, opt, painter, self)

        # Let's be a little economical and only repaint the visible region
        x_start, y_start, x_end, y_end = event.rect().getCoords()
        i_start, j_start = self.coordsToTileIndices(x_start, y_start)
        i_end, j_end = self.coordsToTileIndices(x_end, y_end)

        for i in range(i_start, i_end+1):
            for j in range(j_start, j_end+1):
                tile = self.tiles[i, j]
                if tile:
                    rect = QC.QRect(
                        self.spacing + j*self.tilesize,
                        self.spacing + i*self.tilesize,
                        self.tilesize - self.spacing,
                        self.tilesize - self.spacing
                    )
                    flipx, flipy = self.flips[i, j]
                    painter.drawPixmap(
                        rect,
                        tile.pixmap.transformed(QG.QTransform().scale(flipx, flipy))
                    )

    def coordsToTileIndices(self, x, y):
        j, i = (
            (x - self.spacing)//self.tilesize,
            (y - self.spacing)//self.tilesize,
        )
        return (i, j)

    def resize(self):
        self.tilesize = self.spacing + self.scalefactor*8
        self.setFixedSize(QC.QSize(
            self.tilesize*self.numcols + self.spacing,
            self.tilesize*self.numrows + self.spacing,
        ))
        self.tiles.resize(self.numrows, self.numcols)
        self.update()

    def zoomOut(self):
        self.scalefactor -= 1
        if self.scalefactor < 1:
            self.scalefactor = 1
            return

        self.resize()

    def zoomIn(self):
        self.scalefactor += 1
        if self.scalefactor > 10:
            self.scalefactor = 10
            return

        self.resize()

    def zoomOriginal(self):
        self.scalefactor = self.DEFAULT_SCALE_FACTOR

        self.resize()

    def computeNumRows(self):
        pass

    def increaseCols(self):
        maxcols = len(self.tiles)
        self.numcols += 1
        if self.numcols > maxcols:
            self.numcols = maxcols
        self.computeNumRows()
        self.resize()

    def decreaseCols(self):
        self.numcols -= 1
        if self.numcols < 1:
            self.numcols = 1
        self.computeNumRows()
        self.resize()


class ROMCanvas(TileGrid):
    CSS = """
    background-image: url(img/checkerboard.png);
    background-repeat: repeat-xy;
    """

    def __init__(self, rom):
        self.rom = rom
        tiles = rom.tiles

        super().__init__(
            numcols=16,
            numrows=len(tiles)//16,
            spacing=2,
            scalefactor=5
        )

        self.tiles = GridHolder(tiles, self.numrows, self.numcols)
        self.flips = GridHolder(None, self.numrows, self.numcols, fillvalue=(1, 1))

        self.picked_tile = None

        self.resize()

    def computeNumRows(self):
        self.numrows = len(self.tiles)//self.numcols

    def mousePressEvent(self, event):
        i, j = self.coordsToTileIndices(event.x(), event.y())
        tile = self.tiles[i, j]

        self.picked_tile = tile


class TilePicker(TileGrid):
    CSS = """
    background-image: url(img/grey-checkerboard.png);
    background-repeat: repeat-xy;
    """

    def __init__(self, rom_canvas=None):
        super().__init__(
            numcols=50,
            numrows=30,
            spacing=0,
            scalefactor=5,
        )

        self.rom_canvas = rom_canvas
        self.tiles = GridHolder(None, self.numrows, self.numcols)
        self.flips = GridHolder(None, self.numcols, self.numcols, fillvalue=(1, 1))

        self.resize()

    def mousePressEvent(self, event):
        i, j = self.coordsToTileIndices(event.x(), event.y())

        if event.button() == QC.Qt.LeftButton:
            self.tiles[i, j] = self.rom_canvas.picked_tile
            self.flips[i, j] = (1, 1)
        elif event.button() == QC.Qt.RightButton:
            self.tiles[i, j] = None
        elif event.button() == QC.Qt.MidButton:
            flipx, flipy = self.flips[i, j]
            if flipx == 1:
                if flipy == 1:
                    self.flips[i, j] = (1, -1)
                else:
                    self.flips[i, j] = (-1, 1)
            else:
                if flipy == 1:
                    self.flips[i, j] = (-1, -1)
                else:
                    self.flips[i, j] = (1, 1)

        self.update()


class ColourPicker(QW.QDialog):

    def __init__(self):
        super().__init__()
        self.setWindowTitle("Pick a colour")

        layout = QW.QGridLayout()
        layout.setSpacing(0)

        for i in range(4):
            for j in range(16):
                colour_idx = i*16 + j
                button = ColourButton(colour_idx)
                button.pressed.connect(lambda c=colour_idx: self.colour_picked(c))
                layout.addWidget(button, i, j)
        colours = QW.QWidget()
        colours.setLayout(layout)

        vlayout = QW.QVBoxLayout()
        vlayout.addWidget(colours)

        transparent_button = QW.QPushButton("Transparent")

        transparent_button.pressed.connect(lambda: self.colour_picked(None))
        vlayout.addWidget(transparent_button)
        self.setLayout(vlayout)

    def colour_picked(self, colour_idx):
        self.picked_idx = colour_idx
        self.accept()


class ColourButton(QW.QPushButton):

    def __init__(self, colour_idx):
        super().__init__()
        self.setFixedSize(QC.QSize(32, 32))
        self.colour_idx = colour_idx
        self.set_colour(colour_idx)
        self.setToolTip("Change colour")

    def set_colour(self, colour_idx):
        self.colour_idx = colour_idx

        if colour_idx is None:
            # Enable transparency
            self.setText("")
            self.setStyleSheet("")
            return

        self.setText(f"{colour_idx:0{2}X}")

        bgcolour = NES_PALETTE[colour_idx]
        qt_colour = QG.QColor(bgcolour)

        textcolour = 'white' if is_dark(qt_colour) else 'black'
        bordercolour = "#444444" if qt_colour.rgb() == 0 else qt_colour.darker().name()
        self.setStyleSheet(
            f'''
            QPushButton {{
              color: {textcolour};
              background: {bgcolour};
              border: 2px outset {bordercolour};
              border-radius: 4px;
            }}
            QPushButton:pressed {{
              border-style: inset;
            }}'''
        )


class PaletteButton(ColourButton):

    def __init__(self, button_idx, colour_idx):
        super().__init__(colour_idx)
        self.button_idx = button_idx


class MainWindow(QW.QMainWindow):
    def __init__(self):
        super().__init__()

        bar = self.menuBar()
        file = bar.addMenu("&File")
        open_action = file.addAction("&Open")
        quit_action = file.addAction("&Quit")
        open_action.triggered.connect(self.open_rom)
        quit_action.triggered.connect(app.quit)

        self.tile_picker = TilePicker()

        scroll = QW.QScrollArea()
        scroll.setWidget(self.tile_picker)

        self.setCentralWidget(scroll)

    def open_rom(self):
        filename = QW.QFileDialog.getOpenFileName(
            self,
            "Open ROM", "", "NES Files (*.nes)"
        )
        rom = NES_ROM(filename[0])

        self.grid_widget = ROMCanvas(rom)

        scroll = QW.QScrollArea()
        scroll.setWidget(self.grid_widget)

        self.colour_picker = ColourPicker()

        self.palette = QW.QHBoxLayout()
        for button_idx, color_idx in enumerate(Tile.default_palette):
            button = PaletteButton(button_idx, color_idx)
            button.pressed.connect(lambda idx=button_idx: self.pick_palette_colour(idx))
            self.palette.addWidget(button)

        for zoom in ["in", "out", "original"]:
            zoomButton = QW.QPushButton("")
            iconpath = widget_icon_path(zoomButton)
            zoomButton.setIcon(QG.QIcon(f"img/{iconpath}/zoom-{zoom}.svg"))
            zoomButton.setFixedSize(QC.QSize(48, 48))
            zoomButton.setIconSize(QC.QSize(32, 32))
            zoomButton.setToolTip(f"Zoom {zoom}")
            zoomButton.pressed.connect(getattr(self.grid_widget, f"zoom{zoom.capitalize()}"))
            self.palette.addWidget(zoomButton)

        left = QW.QVBoxLayout()
        left.addWidget(scroll)
        romWidget = QW.QWidget()
        romWidget.setLayout(left)
        paletteWidget = QW.QWidget()
        paletteWidget.setLayout(self.palette)
        left.addWidget(paletteWidget)

        left_dock = QW.QDockWidget("ROM tiles", self)
        left_dock.setWidget(romWidget)

        iconpath = widget_icon_path(left_dock)
        left_dock.setStyleSheet(
           f"""
            QDockWidget
            {{
              titlebar-close-icon: url(img/{iconpath}/widget-close.svg);
              titlebar-normal-icon: url(img/{iconpath}/widget-undock.svg);
            }}
            """
        )

        self.addDockWidget(QC.Qt.LeftDockWidgetArea, left_dock)
        self.tile_picker.rom_canvas = self.grid_widget

    def pick_palette_colour(self, palette_idx):
        if self.colour_picker.exec_():
            new_colour_idx = self.colour_picker.picked_idx
            self.palette.itemAt(palette_idx).widget().set_colour(new_colour_idx)
            for tile in self.grid_widget.tiles:
                palette = tile.palette
                palette[palette_idx] = new_colour_idx
                tile.set_palette(palette)
            self.grid_widget.update()


app = QW.QApplication(sys.argv)
window = MainWindow()
window.show()
app.exec_()