from ....entities import (
    Input,
    TransformerEngineSettings,
    VisionColorModel,
    VisionImageFormat,
)
from ....model.engine import DiffusionPipeline, Engine, PreTrainedModel
from ....model.vendor import TextGenerationVendor
from ....model.vision import BaseVisionModel

from dataclasses import replace
from importlib import import_module
from logging import Logger, getLogger
from types import MethodType
from typing import Any, Literal, cast

from torch import float32, inference_mode


def _diffusion_pipeline() -> Any:
    return getattr(import_module("diffusers"), "DiffusionPipeline")


class TextToImageModel(BaseVisionModel):
    _base: DiffusionPipeline

    def __init__(
        self,
        model_id: str,
        settings: TransformerEngineSettings | None = None,
        logger: Logger = getLogger(__name__),
    ):
        settings = settings or TransformerEngineSettings()
        assert settings.refiner_model_id
        settings = replace(settings, enable_eval=False)
        super().__init__(model_id, settings, logger)

    def _load_model(
        self,
    ) -> PreTrainedModel | TextGenerationVendor | DiffusionPipeline:
        dtype = Engine.weight(self._settings.weight_type)
        dtype_variant = self._settings.weight_type
        pipeline = _diffusion_pipeline()

        base = cast(
            DiffusionPipeline,
            cast(Any, pipeline).from_pretrained(
                self._model_id,
                torch_dtype=dtype,
                variant=dtype_variant,
                use_safetensors=True,
            ),
        )
        cast(Any, base).to(self._device)
        self._replace_deprecated_upcast_vae(base)

        refiner = cast(
            DiffusionPipeline,
            cast(Any, pipeline).from_pretrained(
                self._settings.refiner_model_id,
                text_encoder_2=cast(Any, base).text_encoder_2,
                vae=cast(Any, base).vae,
                torch_dtype=dtype,
                use_safetensors=True,
                variant=dtype_variant,
            ),
        )
        cast(Any, refiner).to(self._device)
        self._replace_deprecated_upcast_vae(refiner)

        self._base = base

        return refiner

    @staticmethod
    def _replace_deprecated_upcast_vae(pipeline: DiffusionPipeline) -> None:
        assert pipeline

        def upcast_vae(self: DiffusionPipeline) -> None:
            cast(Any, self).vae.to(dtype=float32)

        cast(Any, pipeline).upcast_vae = MethodType(upcast_vae, pipeline)

    async def __call__(
        self,
        input: Input,
        path: str,
        *,
        color_model: VisionColorModel = VisionColorModel.RGB,
        high_noise_frac: float = 0.8,
        image_format: VisionImageFormat = VisionImageFormat.JPEG,
        n_steps: int = 150,
        output_type: Literal["latent"] = "latent",
    ) -> str:
        assert (
            input
            and path
            and color_model
            and high_noise_frac is not None
            and image_format
            and n_steps
            and output_type
        )

        with inference_mode():
            image = cast(Any, self._base)(
                prompt=input if isinstance(input, str) else str(input),
                num_inference_steps=n_steps,
                denoising_end=high_noise_frac,
                output_type=output_type,
            ).images
            image = self._model(
                prompt=input if isinstance(input, str) else str(input),
                num_inference_steps=n_steps,
                denoising_start=high_noise_frac,
                image=image,
            ).images[0]

        image.convert(color_model)
        image.save(path, image_format)

        return path
