#!/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.QtGui import QImage, QImageReader

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
