#!/usr/bin/env python3
# upscale-qt - upscale an image with sd-cli and compare before/after
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

import os
import sys
from dataclasses import dataclass
from pathlib import Path

from PyQt6.QtCore import QPointF, QProcess, QRectF, QSettings, QTemporaryDir, Qt
from PyQt6.QtGui import (QAction, QColor, QImage, QImageReader, QKeySequence,
                         QPainter, QPen)
from PyQt6.QtWidgets import (QApplication, QComboBox, QDialog,
                             QDialogButtonBox, QFileDialog, QFormLayout,
                             QLineEdit, QMainWindow, QMessageBox, QWidget)

MODEL_EXTS = ('.pth', '.safetensors', '.gguf')
JPEG_QUALITY = 95
DESCRIPTIONS = 'models.txt'   # optional, in the upscalers folder
DEFAULTS = {'upscalers_dir': '/data/LLM-models/SD/upscalers'}


def list_models(folder):
    try:
        return sorted(f for f in os.listdir(folder)
                      if f.lower().endswith(MODEL_EXTS))
    except OSError:
        return []


def read_descriptions(folder):
    """Parse models.txt: 'filename = what it is for' per line, # comments."""
    try:
        with open(os.path.join(folder, DESCRIPTIONS), encoding='utf-8') as f:
            lines = f.read().splitlines()
    except OSError:
        return {}
    out = {}
    for line in lines:
        name, sep, desc = line.partition('=')
        if sep and not line.lstrip().startswith('#'):
            out[name.strip()] = desc.strip()
    return out


def model_label(name, descriptions):
    desc = descriptions.get(name)
    return f'{desc}  [{Path(name).stem}]' if desc else name


def sd_args(model, src, out):
    return ['-M', 'upscale', '--upscale-model', model, '-i', src, '-o', out]


def save_over(img, path):
    """Replace path with img in the original's format, never truncating it."""
    fmt = bytes(QImageReader.imageFormat(path)).decode() or None
    part = path + '.part'
    if not img.save(part, fmt, JPEG_QUALITY if fmt == 'jpeg' else -1):
        if os.path.exists(part):
            os.remove(part)
        return False
    os.replace(part, path)
    return True


def tail(data, lines=15):
    text = bytes(data).decode(errors='replace').replace('\r', '\n')
    return '\n'.join([l for l in text.splitlines() if l.strip()][-lines:])


@dataclass
class Job:
    path: str          # file the user opened, overwrite target
    image: QImage      # original, EXIF orientation applied
    input: str         # oriented PNG fed to sd-cli
    model: str = ''
    out: str = ''
    result: QImage | None = None
    saved: bool = False


class CompareView(QWidget):
    """Before left of a draggable divider, after right, one shared transform.

    Before is drawn into after's rectangle, so it is scaled on paint, never
    copied at 4x size."""

    GRAB = 8  # px around the divider that start a divider drag

    def __init__(self):
        super().__init__()
        self.before = self.after = None
        self.split = 0.5
        self.scale = 1.0
        self.offset = QPointF()
        self.drag = None
        self.setMouseTracking(True)
        self.setMinimumSize(400, 300)

    def set_images(self, before, after=None):
        self.before, self.after = before, after
        self.fit()

    def ref(self):
        return self.after if self.after is not None else self.before

    def fit(self):
        img = self.ref()
        if img is not None:
            self.scale = min(self.width() / img.width(),
                             self.height() / img.height())
            self.offset = QPointF((self.width() - img.width() * self.scale) / 2,
                                  (self.height() - img.height() * self.scale) / 2)
        self.update()

    def divider_x(self):
        return self.split * self.width()

    def paintEvent(self, e):
        p = QPainter(self)
        p.fillRect(self.rect(), QColor(40, 40, 40))
        img = self.ref()
        if img is None:
            return
        if self.scale >= 0.5:
            # ponytail: smooth only near 1:1, smoothing a 160 MP image
            # scaled down to screen size stalls the paint
            p.setRenderHint(QPainter.RenderHint.SmoothPixmapTransform)
        target = QRectF(self.offset.x(), self.offset.y(),
                        img.width() * self.scale, img.height() * self.scale)
        if self.after is None:
            p.drawImage(target, self.before)
            return
        x = self.divider_x()
        p.setClipRect(QRectF(0, 0, x, self.height()))
        p.drawImage(target, self.before)
        p.setClipRect(QRectF(x, 0, self.width() - x, self.height()))
        p.drawImage(target, self.after)
        p.setClipping(False)
        p.setPen(QPen(QColor('white'), 2))
        p.drawLine(QPointF(x, 0), QPointF(x, self.height()))

    def near_divider(self, pos):
        return self.after is not None and abs(pos.x() - self.divider_x()) < self.GRAB

    def mousePressEvent(self, e):
        self.drag = 'split' if self.near_divider(e.position()) else e.position()

    def mouseMoveEvent(self, e):
        pos = e.position()
        if self.drag == 'split':
            self.split = min(max(pos.x() / self.width(), 0.0), 1.0)
        elif self.drag is not None:
            self.offset += pos - self.drag
            self.drag = pos
        self.setCursor(Qt.CursorShape.SplitHCursor if self.drag == 'split'
                       or self.near_divider(pos) else Qt.CursorShape.ArrowCursor)
        self.update()

    def mouseReleaseEvent(self, e):
        self.drag = None

    def mouseDoubleClickEvent(self, e):
        self.fit()

    def wheelEvent(self, e):
        f = 1.25 if e.angleDelta().y() > 0 else 0.8
        pos = e.position()
        self.offset = pos - (pos - self.offset) * f   # zoom around the cursor
        self.scale *= f
        self.update()

    def resizeEvent(self, e):
        self.fit()  # ponytail: resize resets zoom, keep the view center if it annoys


class SettingsDialog(QDialog):
    def __init__(self, parent, cfg):
        super().__init__(parent)
        self.setWindowTitle('Settings')
        self.cfg = cfg
        self.edits = {k: QLineEdit(str(cfg.value(k, v))) for k, v in DEFAULTS.items()}
        form = QFormLayout(self)
        form.addRow('Upscalers folder', self.edits['upscalers_dir'])
        bb = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok
                              | QDialogButtonBox.StandardButton.Cancel)
        bb.accepted.connect(self.accept)
        bb.rejected.connect(self.reject)
        form.addRow(bb)

    def accept(self):
        for k, e in self.edits.items():
            self.cfg.setValue(k, e.text().strip())
        super().accept()


class Main(QMainWindow):
    def __init__(self, path=None):
        super().__init__()
        self.setWindowTitle('upscale-qt')
        self.resize(1200, 800)
        self.cfg = QSettings('upscale-qt', 'upscale-qt')
        self.tmp = QTemporaryDir()
        self.jobs = []      # one job for now, batch mode adds more
        self.cur = None
        self.runs = 0       # unique temp file names
        self.what = None    # 'upscale' while sd-cli runs
        self.err = b''

        self.view = CompareView()
        self.setCentralWidget(self.view)
        self.models = QComboBox()
        self.models.setSizeAdjustPolicy(QComboBox.SizeAdjustPolicy.AdjustToContents)
        self.models.setToolTip('Upscaler model')

        tb = self.addToolBar('main')
        tb.setMovable(False)

        def act(text, slot, key=None):
            a = QAction(text, self)
            a.triggered.connect(slot)
            if key is not None:
                a.setShortcut(key)
            tb.addAction(a)
            return a

        self.a_open = act('Open', self.open_dialog, QKeySequence.StandardKey.Open)
        tb.addWidget(self.models)
        self.a_up = act('Upscale', self.upscale, 'Ctrl+U')
        self.a_save = act('Save As', self.save_as, QKeySequence.StandardKey.SaveAs)
        self.a_over = act('Overwrite', self.overwrite)
        self.a_clear = act('Clear', self.clear, 'Ctrl+W')
        act('Settings', self.settings)

        self.proc = QProcess(self)
        self.proc.finished.connect(self.proc_done)
        self.proc.errorOccurred.connect(self.proc_error)
        self.proc.readyReadStandardError.connect(self.proc_stderr)

        self.reload_models()
        if path:
            self.load(path)

    def get(self, key):
        return str(self.cfg.value(key, DEFAULTS[key]))

    def job(self):
        return self.jobs[self.cur] if self.cur is not None else None

    def reload_models(self):
        keep = self.models.currentData()
        folder = self.get('upscalers_dir')
        descriptions = read_descriptions(folder)
        self.models.clear()
        for name in list_models(folder):
            self.models.addItem(model_label(name, descriptions), name)
            self.models.setItemData(self.models.count() - 1, name,
                                    Qt.ItemDataRole.ToolTipRole)
        i = self.models.findData(keep)
        if i >= 0:
            self.models.setCurrentIndex(i)
        self.refresh()

    def refresh(self):
        busy = self.what is not None
        job = self.job()
        done = job is not None and job.result is not None
        self.a_open.setEnabled(not busy)
        self.models.setEnabled(not busy)
        self.a_up.setEnabled(job is not None and not busy and self.models.count() > 0)
        self.a_up.setToolTip('Upscale' if self.models.count() else
                             f'No models in {self.get("upscalers_dir")}, see Settings')
        self.a_save.setEnabled(done and not busy)
        self.a_over.setEnabled(done and not busy)
        self.a_clear.setEnabled(job is not None and not busy)

    def open_dialog(self):
        path, _ = QFileDialog.getOpenFileName(
            self, 'Open image', str(self.cfg.value('last_dir', str(Path.home()))),
            'Images (*.png *.jpg *.jpeg *.webp *.bmp *.tif *.tiff);;All files (*)')
        if path:
            self.load(path)

    def load(self, path):
        r = QImageReader(path)
        r.setAutoTransform(True)
        img = r.read()
        if img.isNull():
            QMessageBox.warning(self, 'upscale-qt', f'Cannot read {path}:\n{r.errorString()}')
            return
        self.runs += 1
        src = self.tmp.filePath(f'in{self.runs}.png')
        if not img.save(src):
            QMessageBox.warning(self, 'upscale-qt', f'Cannot write temp file {src}')
            return
        self.jobs = [Job(path, img, src)]
        self.cur = 0
        self.cfg.setValue('last_dir', str(Path(path).parent))
        self.setWindowTitle(f'upscale-qt - {Path(path).name}')
        self.view.set_images(img)
        self.statusBar().showMessage(f'{img.width()}x{img.height()}')
        self.refresh()

    def run(self, what, prog, args, msg):
        self.what, self.err = what, b''
        self.statusBar().showMessage(msg)
        self.refresh()
        self.proc.start(prog, args)

    def upscale(self):
        job = self.job()
        job.model = self.models.currentData()
        self.runs += 1
        job.out = self.tmp.filePath(f'out{self.runs}.png')
        self.run('upscale', 'sd-cli',
                 sd_args(os.path.join(self.get('upscalers_dir'), job.model),
                         job.input, job.out),
                 f'Upscaling with {job.model}...')

    def proc_stderr(self):
        self.err = (self.err + bytes(self.proc.readAllStandardError()))[-8000:]
        last = tail(self.err, 1)
        if last:
            self.statusBar().showMessage(last)

    def proc_done(self, code, status):
        self.what = None
        ok = code == 0 and status == QProcess.ExitStatus.NormalExit
        job = self.job()
        img = QImage(job.out) if ok else QImage()
        if img.isNull():
            self.statusBar().showMessage('Upscale failed')
            QMessageBox.warning(self, 'Upscale failed', tail(self.err) or f'exit code {code}')
        else:
            job.result = img
            self.view.set_images(job.image, img)
            self.statusBar().showMessage(
                f'{job.image.width()}x{job.image.height()} -> '
                f'{img.width()}x{img.height()} with {job.model}')
        self.refresh()

    def proc_error(self, err):
        if err != QProcess.ProcessError.FailedToStart:
            return  # crashes also emit finished, handled there
        self.what = None
        prog = self.proc.program()
        QMessageBox.warning(self, 'upscale-qt', f'Cannot start {prog}: is it installed?')
        self.statusBar().showMessage(f'Cannot start {prog}')
        self.refresh()

    def save_as(self):
        job = self.job()
        src = Path(job.path)
        path, _ = QFileDialog.getSaveFileName(
            self, 'Save upscaled image', str(src.with_name(f'{src.stem}_x4{src.suffix}')),
            'Images (*.png *.jpg *.jpeg *.webp *.tif *.tiff)')
        if not path:
            return
        if job.result.save(path, None, JPEG_QUALITY):
            job.saved = True
            self.statusBar().showMessage(f'Saved {path}')
        else:
            QMessageBox.warning(self, 'upscale-qt', f'Cannot save {path}')

    def overwrite(self):
        job = self.job()
        r = job.result
        if QMessageBox.question(
                self, 'Overwrite',
                f'Replace {job.path}\nwith the {r.width()}x{r.height()} upscaled image?\n'
                'This cannot be undone.') != QMessageBox.StandardButton.Yes:
            return
        if save_over(r, job.path):
            job.saved = True
            self.statusBar().showMessage(f'Overwrote {job.path}')
        else:
            QMessageBox.warning(self, 'upscale-qt', f'Cannot write {job.path}, original untouched')

    def clear(self):
        job = self.job()
        if job.result is not None and not job.saved and QMessageBox.question(
                self, 'Clear', 'The upscaled image is not saved. Discard it?'
                ) != QMessageBox.StandardButton.Yes:
            return
        for f in (job.input, job.out):
            if f and os.path.exists(f):
                os.remove(f)   # 4x PNGs run to hundreds of MB
        self.jobs, self.cur = [], None
        self.view.set_images(None)
        self.setWindowTitle('upscale-qt')
        self.statusBar().clearMessage()
        self.refresh()

    def settings(self):
        if SettingsDialog(self, self.cfg).exec():
            self.reload_models()

    def closeEvent(self, e):
        if self.proc.state() != QProcess.ProcessState.NotRunning:
            self.proc.kill()
            self.proc.waitForFinished(2000)
        super().closeEvent(e)


def main():
    app = QApplication(sys.argv)
    w = Main(sys.argv[1] if len(sys.argv) > 1 else None)
    w.show()
    sys.exit(app.exec())


if __name__ == '__main__':
    main()
