# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import List

import PIL.Image
import torch

from ...utils import logging
from ..modular_pipeline import SequentialPipelineBlocks
from ..modular_pipeline_utils import InsertableDict, OutputParam
from .before_denoise import (
    QwenImageEditPlusRoPEInputsStep,
    QwenImagePrepareLatentsStep,
    QwenImageSetTimestepsStep,
)
from .decoders import (
    QwenImageAfterDenoiseStep,
    QwenImageDecoderStep,
    QwenImageProcessImagesOutputStep,
)
from .denoise import (
    QwenImageEditDenoiseStep,
)
from .encoders import (
    QwenImageEditPlusProcessImagesInputStep,
    QwenImageEditPlusResizeStep,
    QwenImageEditPlusTextEncoderStep,
    QwenImageVaeEncoderStep,
)
from .inputs import (
    QwenImageEditPlusAdditionalInputsStep,
    QwenImageTextInputsStep,
)


logger = logging.get_logger(__name__)


# ====================
# 1. TEXT ENCODER
# ====================


class QwenImageEditPlusVLEncoderStep(SequentialPipelineBlocks):
    """VL encoder that takes both image and text prompts. Uses 384x384 target area."""

    model_name = "qwenimage-edit-plus"
    block_classes = [
        QwenImageEditPlusResizeStep(target_area=384 * 384, output_name="resized_cond_image"),
        QwenImageEditPlusTextEncoderStep(),
    ]
    block_names = ["resize", "encode"]

    @property
    def description(self) -> str:
        return "QwenImage-Edit Plus VL encoder step that encodes the image and text prompts together."


# ====================
# 2. VAE ENCODER
# ====================


class QwenImageEditPlusVaeEncoderStep(SequentialPipelineBlocks):
    """VAE encoder that handles multiple images with different sizes. Uses 1024x1024 target area."""

    model_name = "qwenimage-edit-plus"
    block_classes = [
        QwenImageEditPlusResizeStep(target_area=1024 * 1024, output_name="resized_image"),
        QwenImageEditPlusProcessImagesInputStep(),
        QwenImageVaeEncoderStep(),
    ]
    block_names = ["resize", "preprocess", "encode"]

    @property
    def description(self) -> str:
        return (
            "VAE encoder step that encodes image inputs into latent representations.\n"
            "Each image is resized independently based on its own aspect ratio to 1024x1024 target area."
        )


# ====================
# 3. DENOISE (input -> prepare_latents -> set_timesteps -> prepare_rope_inputs -> denoise -> after_denoise)
# ====================


# assemble input steps
class QwenImageEditPlusInputStep(SequentialPipelineBlocks):
    model_name = "qwenimage-edit-plus"
    block_classes = [
        QwenImageTextInputsStep(),
        QwenImageEditPlusAdditionalInputsStep(image_latent_inputs=["image_latents"]),
    ]
    block_names = ["text_inputs", "additional_inputs"]

    @property
    def description(self):
        return (
            "Input step that prepares the inputs for the Edit Plus denoising step. It:\n"
            " - Standardizes text embeddings batch size.\n"
            " - Processes list of image latents: patchifies, concatenates along dim=1, expands batch.\n"
            " - Outputs lists of image_height/image_width for RoPE calculation.\n"
            " - Defaults height/width from last image in the list."
        )


# Qwen Image Edit Plus (image2image) core denoise step
class QwenImageEditPlusCoreDenoiseStep(SequentialPipelineBlocks):
    model_name = "qwenimage-edit-plus"
    block_classes = [
        QwenImageEditPlusInputStep(),
        QwenImagePrepareLatentsStep(),
        QwenImageSetTimestepsStep(),
        QwenImageEditPlusRoPEInputsStep(),
        QwenImageEditDenoiseStep(),
        QwenImageAfterDenoiseStep(),
    ]
    block_names = [
        "input",
        "prepare_latents",
        "set_timesteps",
        "prepare_rope_inputs",
        "denoise",
        "after_denoise",
    ]

    @property
    def description(self):
        return "Core denoising workflow for QwenImage-Edit Plus edit (img2img) task."

    @property
    def outputs(self):
        return [
            OutputParam(
                name="latents", type_hint=torch.Tensor, description="The latents generated by the denoising step"
            ),
        ]


# ====================
# 4. DECODE
# ====================


class QwenImageEditPlusDecodeStep(SequentialPipelineBlocks):
    model_name = "qwenimage-edit-plus"
    block_classes = [QwenImageDecoderStep(), QwenImageProcessImagesOutputStep()]
    block_names = ["decode", "postprocess"]

    @property
    def description(self):
        return "Decode step that decodes the latents to images and postprocesses the generated image."


# ====================
# 5. AUTO BLOCKS & PRESETS
# ====================

EDIT_PLUS_AUTO_BLOCKS = InsertableDict(
    [
        ("text_encoder", QwenImageEditPlusVLEncoderStep()),
        ("vae_encoder", QwenImageEditPlusVaeEncoderStep()),
        ("denoise", QwenImageEditPlusCoreDenoiseStep()),
        ("decode", QwenImageEditPlusDecodeStep()),
    ]
)


class QwenImageEditPlusAutoBlocks(SequentialPipelineBlocks):
    model_name = "qwenimage-edit-plus"
    block_classes = EDIT_PLUS_AUTO_BLOCKS.values()
    block_names = EDIT_PLUS_AUTO_BLOCKS.keys()

    @property
    def description(self):
        return (
            "Auto Modular pipeline for edit (img2img) tasks using QwenImage-Edit Plus.\n"
            "- `image` is required input (can be single image or list of images).\n"
            "- Each image is resized independently based on its own aspect ratio.\n"
            "- VL encoder uses 384x384 target area, VAE encoder uses 1024x1024 target area."
        )

    @property
    def outputs(self):
        return [
            OutputParam(name="images", type_hint=List[List[PIL.Image.Image]], description="The generated images"),
        ]
