# upscale-qt - tests for the non-GUI helpers and a GUI smoke test
# 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 tempfile
from importlib.machinery import SourceFileLoader
from importlib.util import module_from_spec, spec_from_loader

TMP = tempfile.mkdtemp(prefix='upscale-qt-test-')
os.environ['QT_QPA_PLATFORM'] = 'offscreen'
os.environ['XDG_CONFIG_HOME'] = os.path.join(TMP, 'config')

HERE = os.path.dirname(os.path.abspath(__file__))
_loader = SourceFileLoader('upscale_qt', os.path.join(HERE, 'upscale-qt'))
uq = module_from_spec(spec_from_loader('upscale_qt', _loader))
_loader.exec_module(uq)

from PyQt6.QtGui import QColor, QImage, QImageReader  # noqa: E402

NAMES = ['4x_detail.pth', 'RealESRGAN_x4plus.pth',
         'RealESRGAN_x4plus_anime_6B.pth']


def test_match_model():
    m = uq.match_model
    # the longer stem must win over its own prefix
    assert m('RealESRGAN_x4plus_anime_6B.pth', NAMES) == 'RealESRGAN_x4plus_anime_6B.pth'
    assert m('  realesrgan_x4plus\n', NAMES) == 'RealESRGAN_x4plus.pth'
    assert m('I suggest 4x_detail.', NAMES) == '4x_detail.pth'
    assert m('no idea', NAMES) is None
    assert m('', []) is None


def test_list_models():
    d = os.path.join(TMP, 'models')
    os.makedirs(d)
    for f in ('b.safetensors', 'a.pth', 'notes.txt', 'c.GGUF'):
        open(os.path.join(d, f), 'w').close()
    assert uq.list_models(d) == ['a.pth', 'b.safetensors', 'c.GGUF']
    assert uq.list_models(os.path.join(TMP, 'missing')) == []


def test_args():
    a = uq.sd_args('/m/x.pth', '/in.png', '/out.png')
    assert a == ['-M', 'upscale', '--upscale-model', '/m/x.pth',
                 '-i', '/in.png', '-o', '/out.png']
    a = uq.llama_args('/m.gguf', '/mm.gguf', '/in.png', NAMES)
    assert a[a.index('--image') + 1] == '/in.png'
    prompt = a[a.index('-p') + 1]
    assert all(n in prompt for n in NAMES)
    assert 'noise' in prompt


def test_save_over():
    path = os.path.join(TMP, 'photo.JPG')   # uppercase suffix on purpose
    small = QImage(4, 4, QImage.Format.Format_RGB32)
    small.fill(QColor('red'))
    assert small.save(path, 'jpeg')
    big = QImage(16, 16, QImage.Format.Format_RGB32)
    big.fill(QColor('blue'))
    assert uq.save_over(big, path)
    assert QImage(path).width() == 16
    assert bytes(QImageReader.imageFormat(path)) == b'jpeg'
    assert not os.path.exists(path + '.part')


from PyQt6.QtWidgets import QApplication  # noqa: E402

APP = QApplication.instance() or QApplication([])


def test_compare_view():
    before = QImage(10, 10, QImage.Format.Format_RGB32)
    before.fill(QColor('red'))
    after = QImage(40, 40, QImage.Format.Format_RGB32)
    after.fill(QColor('blue'))
    v = uq.CompareView()
    v.resize(400, 400)   # at or above the minimum size
    v.set_images(before, after)
    shot = v.grab().toImage()
    # left of the divider shows before, right shows after
    assert shot.pixelColor(100, 200) == QColor('red')
    assert shot.pixelColor(300, 200) == QColor('blue')


if __name__ == '__main__':
    for name, fn in list(globals().items()):
        if name.startswith('test_'):
            fn()
            print('ok', name)