#!/usr/bin/env python3
# fcaption: simple image caption editor
# Copyright(c) 2015-2022 by wave++ "Yuri D'Elia" <wavexx@thregr.org>
import os, sys
import argparse
import locale

try:
    from PyQt6 import QtCore, QtGui, QtWidgets
except ImportError:
    from PyQt5 import QtCore, QtGui, QtWidgets

APP_DSC = "fgallery image caption editor"
APP_VER = "1.9.1"
ENCODING = locale.getpreferredencoding()
FILE_EXT = ["jpg", "jpeg", "png", "tif", "tiff"]


class ScaledImage(QtWidgets.QLabel):
    def __init__(self):
        super(ScaledImage, self).__init__()
        self._pixmap = QtGui.QPixmap()

    def setPixmap(self, pixmap):
        self._pixmap = pixmap
        if not pixmap.isNull():
            pixmap = pixmap.scaled(self.size(),
                                   QtCore.Qt.AspectRatioMode.KeepAspectRatio,
                                   QtCore.Qt.TransformationMode.SmoothTransformation)
        super(ScaledImage, self).setPixmap(pixmap)

    def resizeEvent(self, ev):
        super(ScaledImage, self).resizeEvent(ev)
        if not self._pixmap.isNull():
            pixmap = self._pixmap.scaled(self.size(),
                                         QtCore.Qt.AspectRatioMode.KeepAspectRatio,
                                         QtCore.Qt.TransformationMode.SmoothTransformation)
            super(ScaledImage, self).setPixmap(pixmap)


class ThumbSignals(QtCore.QObject):
    ready = QtCore.pyqtSignal(int, QtGui.QImage)

class ThumbLoader(QtCore.QRunnable):
    def __init__(self, idx, path, size):
        super(ThumbLoader, self).__init__()
        self.idx = idx
        self.path = path
        self.size = size
        self.signals = ThumbSignals()

    def run(self):
        image = QtGui.QImage(self.path)
        if not image.isNull():
            scaled = image.scaled(
                self.size,
                QtCore.Qt.AspectRatioMode.KeepAspectRatio,
                QtCore.Qt.TransformationMode.SmoothTransformation)
            self.signals.ready.emit(self.idx, scaled)


class MainWindow(QtWidgets.QMainWindow):
    def __init__(self):
        super(MainWindow, self).__init__()
        self.setWindowTitle(APP_DSC)

        # construct UI
        horizontalSplitter = QtWidgets.QSplitter(QtCore.Qt.Orientation.Horizontal)
        verticalSplitter = QtWidgets.QSplitter(QtCore.Qt.Orientation.Vertical)
        self.image = ScaledImage()
        self.image.setMinimumSize(480, 319)
        sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Policy.Expanding,
                                           QtWidgets.QSizePolicy.Policy.Expanding)
        sizePolicy.setHorizontalStretch(1)
        sizePolicy.setVerticalStretch(1)
        self.image.setSizePolicy(sizePolicy)
        self.image.setAlignment(QtCore.Qt.AlignmentFlag.AlignCenter)
        verticalSplitter.addWidget(self.image)
        horizontalLayout = QtWidgets.QHBoxLayout()
        formLayout = QtWidgets.QFormLayout()
        formLayout.setLabelAlignment(QtCore.Qt.AlignmentFlag.AlignRight |
                                     QtCore.Qt.AlignmentFlag.AlignTrailing |
                                     QtCore.Qt.AlignmentFlag.AlignVCenter)
        formLayout.setWidget(0, QtWidgets.QFormLayout.ItemRole.LabelRole, QtWidgets.QLabel("Title:"))
        self.edit_title = QtWidgets.QLineEdit()
        formLayout.setWidget(0, QtWidgets.QFormLayout.ItemRole.FieldRole, self.edit_title)
        formLayout.setWidget(1, QtWidgets.QFormLayout.ItemRole.LabelRole, QtWidgets.QLabel("Description:"))
        self.edit_desc = QtWidgets.QPlainTextEdit()
        self.edit_desc.setLineWrapMode(QtWidgets.QPlainTextEdit.LineWrapMode.WidgetWidth)
        sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Policy.Expanding,
                                           QtWidgets.QSizePolicy.Policy.Preferred)
        sizePolicy.setHorizontalStretch(1)
        sizePolicy.setVerticalStretch(1)
        self.edit_desc.setSizePolicy(sizePolicy)
        formLayout.setWidget(1, QtWidgets.QFormLayout.ItemRole.FieldRole, self.edit_desc)
        horizontalLayout.addLayout(formLayout)
        verticalLayout = QtWidgets.QVBoxLayout()
        verticalLayout.setSizeConstraint(QtWidgets.QLayout.SizeConstraint.SetMinimumSize)
        self.btn_next = QtWidgets.QPushButton("&Next")
        verticalLayout.addWidget(self.btn_next)
        self.btn_undo = QtWidgets.QPushButton("&Undo")
        verticalLayout.addWidget(self.btn_undo)
        self.btn_prev = QtWidgets.QPushButton("&Previous")
        verticalLayout.addWidget(self.btn_prev)
        horizontalLayout.addLayout(verticalLayout)
        widget = QtWidgets.QWidget()
        widget.setLayout(horizontalLayout)
        verticalSplitter.addWidget(widget)
        horizontalSplitter.addWidget(verticalSplitter)
        self.list_files = QtWidgets.QListWidget()
        self.list_files.setIconSize(QtCore.QSize(64, 64))
        horizontalSplitter.addWidget(self.list_files)
        self.setCentralWidget(horizontalSplitter)

        # signals
        self.list_files.itemActivated.connect(self.on_list)
        self.btn_next.clicked.connect(self.on_next)
        self.btn_prev.clicked.connect(self.on_prev)
        self.btn_undo.clicked.connect(self.on_undo)
        self.edit_title.textEdited.connect(self.on_changed)
        self.edit_desc.textChanged.connect(self.on_changed)

        # initial state
        self.files = list()
        self.current = 0
        self.modified = False
        self.thumb_pool = QtCore.QThreadPool()
        self.thumb_pool.setMaxThreadCount(max(1, self.thumb_pool.maxThreadCount() // 4))

    def on_next(self, ev):
        if self.modified: self.save()
        self.load((self.current + 1) % len(self.files))

    def on_prev(self, ev):
        if self.modified: self.save()
        self.load((self.current - 1) % len(self.files))

    def on_list(self, ev):
        if self.modified: self.save()
        self.load(self.list_files.currentRow())

    def on_undo(self, ev):
        self.load(self.current)

    def thumb_ready(self, idx, image):
        crow = self.list_files.currentRow()
        item = self.list_files.takeItem(idx)
        item.setIcon(QtGui.QIcon(QtGui.QPixmap(image)))
        self.list_files.insertItem(idx, item)
        self.list_files.setCurrentRow(crow)

    def thumb_schedule(self):
        size = self.list_files.iconSize()
        for idx, path in enumerate(self.files):
            thread = ThumbLoader(idx, path, size)
            thread.signals.ready.connect(self.thumb_ready)
            self.thumb_pool.start(thread)

    def thumb_stop(self):
        self.thumb_pool.clear()
        self.thumb_pool.waitForDone()

    def set_files(self, files):
        self.files = list(files)
        self.list_files.clear()
        for path in files:
            self.list_files.addItem(os.path.basename(path))
        if len(files) < 2:
            self.list_files.hide()
            self.btn_next.setEnabled(False)
            self.btn_prev.setEnabled(False)
        else:
            self.list_files.show()
            self.thumb_schedule()
            self.btn_next.setEnabled(True)
            self.btn_prev.setEnabled(True)
        self.load(0)

    def on_changed(self, *_):
        self.modified = True

    def load(self, idx):
        self.current = idx
        self.list_files.setCurrentRow(idx)

        path = self.files[idx]
        pixmap = QtGui.QPixmap(path)
        if pixmap.isNull():
            self.image.setPixmap(pixmap)
            self.image.setText('Cannot load {}'.format(path))
            self.edit_title.setEnabled(False)
            self.edit_desc.setEnabled(False)
            return

        self.image.clear()
        self.image.setPixmap(pixmap)

        self.edit_title.clear()
        self.edit_desc.clear()
        self.edit_title.setEnabled(True)
        self.edit_desc.setEnabled(True)

        base, _ = os.path.splitext(path)
        txt = base + '.txt'
        if os.path.isfile(txt):
            data = open(txt, 'rb').read().decode(ENCODING).split('\n', 1)
            if len(data) > 0:
                self.edit_title.setText(data[0].strip())
            if len(data) > 1:
                self.edit_desc.setPlainText(data[1].strip())

        self.modified = False
        self.edit_title.setFocus()

    def save(self):
        title = str(self.edit_title.text()).strip()
        desc = str(self.edit_desc.toPlainText()).strip()
        base, _ = os.path.splitext(self.files[self.current])
        txt = base + '.txt'
        if len(title) + len(desc) == 0:
            if os.path.isfile(txt):
                os.remove(txt)
        else:
            data = title + '\n' + desc
            open(txt, 'wb').write(data.encode(ENCODING))

    def closeEvent(self, ev):
        self.thumb_stop()
        if self.modified: self.save()
        super(MainWindow, self).closeEvent(ev)


# main application
def expand_dir(path):
    for root, dirs, files in os.walk(path):
        files.sort()
        for tmp in files:
            if tmp[0] == '.': continue
            tmp = os.path.join(root, tmp)
            ext = os.path.splitext(tmp)[1]
            if ext: ext = ext[1:].lower()
            if ext in FILE_EXT:
                yield tmp


class Application(QtWidgets.QApplication):
    def __init__(self, args):
        super(Application, self).__init__(args)

        # command-line flags
        ap = argparse.ArgumentParser(description=APP_DSC)
        ap.add_argument('--version', action='version', version=APP_VER)
        ap.add_argument('files', metavar="image", nargs='*',
                        help='image or directory to caption')
        args = ap.parse_args(map(str, args[1:]))

        # ask for a directory if no files were specified
        if not args.files:
            path = QtWidgets.QFileDialog.getExistingDirectory(
                None, "Select an image directory")
            if not path: sys.exit(1)
            args.files = [str(path)]

        # expand directories to files
        files = []
        for path in args.files:
            if not os.path.isdir(path):
                files.append(path)
            else:
                files.extend(expand_dir(path))
        if not files:
            print("no files to caption", file=sys.stderr)
            sys.exit(1)

        # initialize
        self.main_window = MainWindow()
        self.main_window.set_files(files)
        self.main_window.show()

if __name__ == '__main__':
    app = Application(sys.argv)
    sys.exit(app.exec())
