import logging
from typing import TYPE_CHECKING, Any, Dict, List, Sequence

import wandb
from wandb.sdk.integration_utils.auto_logging import Response
from wandb.util import get_module

from .utils import chunkify, get_updated_kwargs

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    torch_float_tensor = get_module("torch.FloatTensor")


SUPPORTED_SDXL_PIPELINES = [
    "StableDiffusionXLPipeline",
    "StableDiffusionXLImg2ImgPipeline",
]

TEXT_TO_IMAGE_COLUMNS = [
    "Workflow-Stage",
    "Prompt",
    "Negative-Prompt",
    "Prompt-2",
    "Negative-Prompt-2",
    "Generated-Image",
]


def decode_sdxl_t2i_latents(pipeline: Any, latents: "torch_float_tensor") -> List:
    """Decode latents generated by [`diffusers.StableDiffusionXLPipeline`](https://huggingface.co/docs/diffusers/main/en/api/pipelines/stable_diffusion/stable_diffusion_xl#stable-diffusion-xl).

    Arguments:
        pipeline: (diffusers.DiffusionPipeline) The Diffusion Pipeline from
            [`diffusers`](https://huggingface.co/docs/diffusers).
        latents (torch.FloatTensor): The generated latents.

    Returns:
        List of `PIL` images corresponding to the generated latents.
    """
    torch = get_module(
        "torch",
        required="Please ensure PyTorch is installed. You can check out https://pytorch.org/get-started/locally/#start-locally for installation instructions.",
    )
    with torch.no_grad():
        needs_upcasting = (
            pipeline.vae.dtype == torch.float16 and pipeline.vae.config.force_upcast
        )
        if needs_upcasting:
            pipeline.upcast_vae()
            latents = latents.to(
                next(iter(pipeline.vae.post_quant_conv.parameters())).dtype
            )
        images = pipeline.vae.decode(
            latents / pipeline.vae.config.scaling_factor, return_dict=False
        )[0]
        if needs_upcasting:
            pipeline.vae.to(dtype=torch.float16)
        if pipeline.watermark is not None:
            images = pipeline.watermark.apply_watermark(images)
        images = pipeline.image_processor.postprocess(images, output_type="pil")
        pipeline.maybe_free_model_hooks()
        return images


class SDXLResolver:
    """Resolver for  request and responses from [`diffusers.StableDiffusionXLPipeline`](https://huggingface.co/docs/diffusers/main/en/api/pipelines/stable_diffusion/stable_diffusion_xl#stable-diffusion-xl) multi-modal Diffusion Pipelines, providing necessary data transformations, formatting, and logging.

    Arguments:
        pipeline_name: (str) The name of the Diffusion Pipeline.
    """

    def __init__(self, pipeline_name: str) -> None:
        self.pipeline_name = pipeline_name
        self.wandb_table = None
        self.task = None
        self.workflow_stage = None

    def __call__(
        self,
        args: Sequence[Any],
        kwargs: Dict[str, Any],
        response: Response,
        start_time: float,
        time_elapsed: float,
    ) -> Any:
        """Main call method for the `DiffusersPipelineResolver` class.

        Arguments:
            args: (Sequence[Any]) List of arguments.
            kwargs: (Dict[str, Any]) Dictionary of keyword arguments.
            response: (wandb.sdk.integration_utils.auto_logging.Response) The response from
                the request.
            start_time: (float) Time when request started.
            time_elapsed: (float) Time elapsed for the request.

        Returns:
            Packed data as a dictionary for logging to wandb, None if an exception occurred.
        """
        try:
            # Get the pipeline and the args
            pipeline, args = args[0], args[1:]

            # Update the Kwargs so that they can be logged easily
            kwargs = get_updated_kwargs(pipeline, args, kwargs)

            # Get the pipeline configs
            pipeline_configs = dict(pipeline.config)
            pipeline_configs["pipeline-name"] = self.pipeline_name

            # Return the WandB loggable dict
            loggable_dict = self.prepare_loggable_dict(
                pipeline, pipeline_configs, response, kwargs
            )
            return loggable_dict
        except Exception as e:
            logger.warning(e)
        return None

    def create_wandb_table(self, pipeline_configs: Dict[str, Any]) -> None:
        """Create the `wandb.Table` with the specified schema.

        Arguments:
            pipeline_configs: (Dict[str, Any]) The configs corresponding to the
                architecture of the pipeline.
        """
        columns = []
        if self.pipeline_name == "StableDiffusionXLPipeline":
            columns += TEXT_TO_IMAGE_COLUMNS
            self.task = "text_to_image"
            self.workflow_stage = "Base"
        elif self.pipeline_name == "StableDiffusionXLImg2ImgPipeline":
            if (
                pipeline_configs["_name_or_path"]
                == "stabilityai/stable-diffusion-xl-refiner-1.0"
            ):
                columns += TEXT_TO_IMAGE_COLUMNS
                self.task = "text_to_image"
                self.workflow_stage = "Refiner"
        self.wandb_table = wandb.Table(columns=columns)

    def prepare_loggable_dict_for_text_to_image(
        self,
        pipeline: Any,
        workflow_stage: str,
        response: Response,
        kwargs: Dict[str, Any],
    ) -> Dict[str, Any]:
        """Prepare the loggable dictionary for text-to-image workflow.

        Arguments:
            pipeline: (Any) The Diffusion Pipeline.
            workflow_stage: (str) The name of the workflow stage.
            response: (wandb.sdk.integration_utils.auto_logging.Response) The response from
                the request.
            kwargs: (Dict[str, Any]) Dictionary of keyword arguments.

        """
        prompt_logging = (
            kwargs["prompt"]
            if isinstance(kwargs["prompt"], list)
            else [kwargs["prompt"]]
        )
        prompt2_logging = (
            kwargs["prompt_2"]
            if isinstance(kwargs["prompt_2"], list)
            else [kwargs["prompt_2"]]
        )
        negative_prompt_logging = (
            kwargs["negative_prompt"]
            if isinstance(kwargs["negative_prompt"], list)
            else [kwargs["negative_prompt"]]
        )
        negative_prompt2_logging = (
            kwargs["negative_prompt_2"]
            if isinstance(kwargs["negative_prompt_2"], list)
            else [kwargs["negative_prompt_2"]]
        )
        images = (
            decode_sdxl_t2i_latents(pipeline, response.images)
            if kwargs["output_type"] == "latent"
            else response.images
        )
        images = chunkify(images, len(prompt_logging))
        for idx in range(len(prompt_logging)):
            for image in images[idx]:
                wandb.log(
                    {
                        f"Generated-Image/{workflow_stage}": wandb.Image(
                            image,
                            caption=f"Prompt-1: {prompt_logging[idx]}\nPrompt-2: {prompt2_logging[idx]}",
                        )
                    }
                )
                self.wandb_table.add_data(
                    workflow_stage,
                    prompt_logging[idx] if prompt_logging[idx] is not None else "",
                    negative_prompt_logging[idx]
                    if negative_prompt_logging[idx] is not None
                    else "",
                    prompt2_logging[idx] if prompt2_logging[idx] is not None else "",
                    negative_prompt2_logging[idx]
                    if negative_prompt2_logging[idx] is not None
                    else "",
                    wandb.Image(image),
                )

    def update_wandb_configs(
        self, pipeline_configs: Dict[str, Any], kwargs: Dict[str, Any]
    ) -> None:
        """Update the configs of the W&B run.

        Arguments:
            pipeline_configs: (Dict[str, Any]) The configs corresponding to the
                architecture of the pipeline.
            kwargs: (Dict[str, Any]) Dictionary of keyword arguments.

        """
        if "workflow" not in wandb.config:
            wandb.config.update(
                {
                    "workflow": [
                        {
                            "pipeline": pipeline_configs,
                            "params": kwargs,
                            "stage": self.workflow_stage,
                        }
                    ]
                }
            )
        else:
            existing_workflow = wandb.config.workflow
            updated_workflow = existing_workflow + [
                {
                    "pipeline": pipeline_configs,
                    "params": kwargs,
                    "stage": self.workflow_stage,
                }
            ]
            wandb.config.update(
                {
                    "workflow": updated_workflow,
                },
                allow_val_change=True,
            )

    def prepare_loggable_dict(
        self,
        pipeline: Any,
        pipeline_configs: Dict[str, Any],
        response: Response,
        kwargs: Dict[str, Any],
    ) -> Dict[str, Any]:
        """Prepare the loggable dictionary, which is the packed data as a dictionary for logging to wandb, None if an exception occurred.

        Arguments:
            pipeline: (Any) The Diffusion Pipeline.
            pipeline_configs: (Dict[str, Any]) The configs corresponding to the
                architecture of the pipeline.
            response: (wandb.sdk.integration_utils.auto_logging.Response) The
                response from the request.
            kwargs: (Dict[str, Any]) Dictionary of keyword arguments.

        Returns:
            Packed data as a dictionary for logging to wandb, None if an exception occurred.
        """
        self.create_wandb_table(pipeline_configs)
        if self.task == "text_to_image":
            self.prepare_loggable_dict_for_text_to_image(
                pipeline, self.workflow_stage, response, kwargs
            )
        self.update_wandb_configs(pipeline_configs, kwargs)
        return {f"{self.workflow_stage}/Result-Table": self.wandb_table}
