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

LORA_MODELS = [
    "p-image-lora",
]


class PrunaTextToImageLoRA:
    """Generate an image from a text prompt with 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-lora must be trained with p-image-trainer.
    """

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

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prompt": ("STRING", {"multiline": True, "default": ""}),
                "model": (LORA_MODELS, {"default": "p-image-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": (ASPECT_RATIOS, {"default": "1:1"}),
                "api_key": ("STRING", {"default": ""}),
            },
            "optional": {
                "hf_api_token": ("STRING", {"default": ""}),
            },
        }

    def generate(
        self,
        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)
        payload: dict = {
            "input": {
                "prompt": prompt,
                "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,)