import torch
from .utils import (
    call_api_with_fallback,
    generation_url_to_tensor,
    resolve_api_key,
    upload_image,
)


class PrunaImageUpscale:
    """Upscale an image using the Pruna API (p-image-upscale).

    Pricing: $0.005/image for 1-4 MP target, $0.01/image for 5-8 MP target.
    """

    CATEGORY = "pruna ai"
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "upscale"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "target_mp": (
                    "INT",
                    {"default": 4, "min": 1, "max": 8, "step": 1},
                ),
                "enhance_details": ("BOOLEAN", {"default": False}),
                "enhance_realism": ("BOOLEAN", {"default": False}),
                "api_key": ("STRING", {"default": ""}),
            }
        }

    def upscale(
        self,
        image: torch.Tensor,
        target_mp: int,
        enhance_details: bool,
        enhance_realism: bool,
        api_key: str,
    ):
        key = resolve_api_key(api_key)
        image_url = upload_image(image, key)
        payload = {
            "input": {
                "image": image_url,
                "target": target_mp,
                "enhance_details": enhance_details,
                "enhance_realism": enhance_realism,
                "output_format": "png",
            }
        }
        data = call_api_with_fallback("p-image-upscale", payload, key)
        tensor = generation_url_to_tensor(data)
        return (tensor,)