from ..media.resolution import DEFAULT_MAX_AI_INPUT_SIZE, limit_image


def _source_alpha(image):
    if "A" not in image.getbands() and "transparency" not in image.info:
        return None
    return image.convert("RGBA").getchannel("A")


def infer_one(
    model_manager,
    parameters,
    image_path,
    cancel_check=None,
    preserve_alpha=False,
    release_memory=True,
):
    from PIL import Image

    from .onnx_upscale import UpscaleResourceError, infer

    model_directory = model_manager.directory(parameters["model"])
    model = model_manager.get_upscale(
        model_directory,
        parameters["device"],
    )
    try:
        with Image.open(image_path) as opened:
            source = limit_image(
                opened,
                parameters.get("max_input_size", DEFAULT_MAX_AI_INPUT_SIZE),
            )
            source_alpha = _source_alpha(source) if preserve_alpha else None
            prediction = infer(
                model,
                source,
                cancel_check=cancel_check,
                release_memory=release_memory,
            )
            if source_alpha is None:
                return prediction
            alpha = source_alpha.resize(
                prediction.size,
                Image.Resampling.LANCZOS,
            )
            prediction = prediction.convert("RGBA")
            prediction.putalpha(alpha)
            return prediction
    except (MemoryError, UnicodeDecodeError, UpscaleResourceError) as error:
        model_manager.release_upscale()
        device = str(parameters.get("device", "selected device")).upper()
        raise RuntimeError(
            f"Upscale could not finish on {device}: {error}. "
            "Reduce Maximum AI Input Size or select CPU, then try again."
        ) from error