# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from __future__ import annotations

import logging
import os

from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.models.dmd2.config import DMD2Config
from vllm_omni.diffusion.models.schedulers import DMD2EulerScheduler
from vllm_omni.diffusion.models.utils import _load_json
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch

logger = logging.getLogger(__name__)


class DMD2PipelineMixin:
    """Mixin for FastGen DMD2-distilled models. Must appear before the base pipeline in MRO."""

    def __init_dmd2__(self) -> None:
        """Call after super().__init__() to apply DMD2 scheduler and read model_index."""
        local_files_only = os.path.exists(self.od_config.model)
        try:
            model_index = _load_json(self.od_config.model, "model_index.json", local_files_only)
        except Exception:
            model_index = {}

        self.dmd2_config = DMD2Config.from_model_index(model_index)

        self.scheduler = DMD2EulerScheduler(
            num_train_timesteps=1000,
            shift=1.0,
            dmd2_timesteps=self.dmd2_config.resolve_timesteps(),
            stochastic_sampling=(self.dmd2_config.solver == "sde"),
        )

    def _sanitize_dmd2_request(self, req) -> None:
        """Sanitize CFG-related fields in-place. Works with both OmniDiffusionRequest and DiffusionRequestBatch."""
        sampling_params = req.sampling_params_list if isinstance(req, DiffusionRequestBatch) else [req.sampling_params]
        for sp in sampling_params:
            if sp.num_inference_steps and sp.num_inference_steps != self.dmd2_config.num_inference_steps:
                logger.warning(
                    "DMD2: ignoring num_inference_steps=%d, forcing %d.",
                    sp.num_inference_steps,
                    self.dmd2_config.num_inference_steps,
                )
            sp.num_inference_steps = self.dmd2_config.num_inference_steps

            if sp.guidance_scale_provided and sp.guidance_scale != self.dmd2_config.guidance_scale:
                logger.warning(
                    "DMD2: ignoring guidance_scale=%.2f, forcing %.2f.",
                    sp.guidance_scale,
                    self.dmd2_config.guidance_scale,
                )
            sp.guidance_scale = self.dmd2_config.guidance_scale
            sp.guidance_scale_provided = False

            if sp.guidance_scale_2 is not None:
                logger.warning("DMD2: ignoring guidance_scale_2.")
                sp.guidance_scale_2 = None

            if sp.true_cfg_scale is not None:
                logger.warning("DMD2: ignoring true_cfg_scale.")
                sp.true_cfg_scale = None

            sp.do_classifier_free_guidance = False
            sp.is_cfg_negative = False

            # Prevent base pipelines from replacing the DMD2 scheduler.
            extra_args = getattr(sp, "extra_args", None) or {}
            for key in ("sample_solver", "flow_shift"):
                if key in extra_args:
                    logger.warning("DMD2: ignoring extra_args.%s.", key)
                    extra_args.pop(key)

        # Strip negative_prompt from each request's prompt in-place.
        requests = req.requests if hasattr(req, "requests") else [req]
        for request in requests:
            p = request.prompt
            if isinstance(p, dict) and "negative_prompt" in p:
                logger.warning("DMD2: ignoring negative_prompt.")
                request.prompt = {k: v for k, v in p.items() if k != "negative_prompt"}

    def forward(self, req: DiffusionRequestBatch, **kwargs) -> list[DiffusionOutput]:
        self._sanitize_dmd2_request(req)
        kwargs.pop("guidance_scale", None)
        kwargs.pop("num_inference_steps", None)
        return super().forward(
            req,
            guidance_scale=self.dmd2_config.guidance_scale,
            num_inference_steps=self.dmd2_config.num_inference_steps,
            **kwargs,
        )