from typing import Any, Dict, Sequence

from wandb.sdk.integration_utils.auto_logging import Response

from .resolvers import (
    SUPPORTED_MULTIMODAL_PIPELINES,
    DiffusersMultiModalPipelineResolver,
)


class DiffusersPipelineResolver:
    """Resolver for `DiffusionPipeline` request and responses from [HuggingFace Diffusers](https://huggingface.co/docs/diffusers/index), providing necessary data transformations, formatting, and logging.

    This is based off `wandb.sdk.integration_utils.auto_logging.RequestResponseResolver`.
    """

    def __init__(self) -> None:
        self.wandb_table = None
        self.pipeline_call_count = 1

    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.

        Args:
            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.
        """
        pipeline_name = args[0].__class__.__name__
        resolver = None
        if pipeline_name in SUPPORTED_MULTIMODAL_PIPELINES:
            resolver = DiffusersMultiModalPipelineResolver(
                pipeline_name, self.pipeline_call_count
            )
            self.pipeline_call_count += 1
        loggable_dict = resolver(args, kwargs, response, start_time, time_elapsed)
        return loggable_dict