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

EDIT_MODELS = [
    "p-image-edit",
]

EDIT_ASPECT_RATIOS = ["match_input_image"] + ASPECT_RATIOS

EDIT_LORA_MODELS = [
    "p-image-edit-lora",
]


class PrunaImageEdit:
    """Edit an image with a text instruction using the Pruna API.

    Accepts a single ComfyUI IMAGE. The image is uploaded to the Pruna files
    API and passed as a URL to the editing model.
    """

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

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "prompt": ("STRING", {"multiline": True, "default": ""}),
                "model": (EDIT_MODELS, {"default": "p-image-edit"}),
                "aspect_ratio": (EDIT_ASPECT_RATIOS, {"default": "match_input_image"}),
                "api_key": ("STRING", {"default": ""}),
            }
        }

    def edit(
        self,
        image: torch.Tensor,
        prompt: str,
        model: str,
        aspect_ratio: str,
        api_key: str,
    ):
        key = resolve_api_key(api_key)
        image_url = upload_image(image, key)
        payload = {
            "input": {
                "prompt": prompt,
                "images": [image_url],
                "aspect_ratio": aspect_ratio,
            }
        }
        data = call_api_with_fallback(model, payload, key)
        tensor = generation_url_to_tensor(data)
        return (tensor,)


class PrunaImageEditLoRA:
    """Edit an image with a text instruction and a custom LoRA using the Pruna API.

    LoRA weights must be a HuggingFace URL in the format:
    huggingface.co/<username>/<repo>[/<filename>]

    Note: LoRAs used with p-image-edit-lora must be trained with p-image-edit-trainer.
    """

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

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "prompt": ("STRING", {"multiline": True, "default": ""}),
                "model": (EDIT_LORA_MODELS, {"default": "p-image-edit-lora"}),
                "lora_weights": (
                    "STRING",
                    {"default": "huggingface.co/username/repo"},
                ),
                "lora_scale": (
                    "FLOAT",
                    {"default": 0.5, "min": -1.0, "max": 3.0, "step": 0.05},
                ),
                "aspect_ratio": (EDIT_ASPECT_RATIOS, {"default": "match_input_image"}),
                "api_key": ("STRING", {"default": ""}),
            },
            "optional": {
                "hf_api_token": ("STRING", {"default": ""}),
            },
        }

    def edit(
        self,
        image: torch.Tensor,
        prompt: str,
        model: str,
        lora_weights: str,
        lora_scale: float,
        aspect_ratio: str,
        api_key: str,
        hf_api_token: str = "",
    ):
        key = resolve_api_key(api_key)
        image_url = upload_image(image, key)
        payload: dict = {
            "input": {
                "prompt": prompt,
                "images": [image_url],
                "lora_weights": lora_weights,
                "lora_scale": lora_scale,
                "aspect_ratio": aspect_ratio,
            }
        }
        add_hf_token(payload, hf_api_token)

        data = call_api_with_fallback(model, payload, key)
        tensor = generation_url_to_tensor(data)
        return (tensor,)