#!/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 shutil
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
DEFAULTS = {
    'upscalers_dir': '/data/LLM-models/SD/upscalers',
    'imajev_model': '/data/LLM-models/SD/imajev/imajev-4b-Q4_K_M.gguf',
    'imajev_mmproj': '/data/LLM-models/SD/imajev/imajev-4b-mmproj-f16.gguf',
}
PROMPT = (
    'You pick an image upscaler model. Judge the subject (photo of people, '
    'landscape, anime, illustration) and also the source quality: noise, '
    'JPEG compression, blur. Noisy or compressed sources need a denoising '
    'model, clean sources a detail-preserving one. Answer with exactly one '
    'filename from this list and nothing else: {names}')


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


def match_model(answer, names):
    """Map imajev's free-text answer to one of names, or None.

    Longest stem first, so 'x4plus_anime_6B' is not taken for 'x4plus'."""
    answer = answer.lower()
    for name in sorted(names, key=lambda n: len(Path(n).stem), reverse=True):
        if Path(name).stem.lower() in answer:
            return name
    return None


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


def llama_args(model, mmproj, src, names):
    return ['-m', model, '--mmproj', mmproj, '--image', src,
            '--temp', '0', '-n', '64',
            '-p', PROMPT.format(names=' '.join(names))]


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 and imajev
    model: str = ''
    out: str = ''
    result: QImage | None = None


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'])
        form.addRow('imajev model (.gguf)', self.edits['imajev_model'])
        form.addRow('imajev mmproj (.gguf)', self.edits['imajev_mmproj'])
        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    # 'suggest' or 'upscale' while a process 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)
        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 names(self):
        return [self.models.itemText(i) for i in range(self.models.count())]

    def reload_models(self):
        keep = self.models.currentText()
        self.models.clear()
        self.models.addItems(list_models(self.get('upscalers_dir')))
        if keep:
            self.models.setCurrentText(keep)
        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)

    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()
        self.suggest()

    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 suggest(self):
        model, mmproj = self.get('imajev_model'), self.get('imajev_mmproj')
        if not (self.names() and os.path.isfile(model) and os.path.isfile(mmproj)
                and shutil.which('llama-mtmd-cli')):
            return  # suggestion is optional, skip quietly
        self.run('suggest', 'llama-mtmd-cli',
                 llama_args(model, mmproj, self.job().input, self.names()),
                 'imajev is looking at the image...')

    def upscale(self):
        job = self.job()
        job.model = self.models.currentText()
        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:]
        if self.what == 'upscale':
            last = tail(self.err, 1)
            if last:
                self.statusBar().showMessage(last)

    def proc_done(self, code, status):
        what, self.what = self.what, None
        ok = code == 0 and status == QProcess.ExitStatus.NormalExit
        job = self.job()
        if what == 'suggest':
            answer = bytes(self.proc.readAllStandardOutput()).decode(errors='replace')
            pick = match_model(answer, self.names()) if ok else None
            if pick:
                self.models.setCurrentText(pick)
                self.statusBar().showMessage(f'imajev suggests {pick}')
            else:
                self.statusBar().showMessage('imajev gave no usable suggestion')
        elif what == 'upscale':
            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
        what, self.what = self.what, None
        prog = self.proc.program()
        if what == 'upscale':
            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):
            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):
            self.statusBar().showMessage(f'Overwrote {job.path}')
        else:
            QMessageBox.warning(self, 'upscale-qt', f'Cannot write {job.path}, original untouched')

    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()
