#!/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 pathlib import Path

from PyQt6.QtCore import QPointF, QRectF, Qt
from PyQt6.QtGui import QColor, QImage, QImageReader, QPainter, QPen
from PyQt6.QtWidgets import QWidget

MODEL_EXTS = ('.pth', '.safetensors', '.gguf')
JPEG_QUALITY = 95
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


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
