"""Reusable utility for building tiny diffusion models with random weights.

Mirrors the component loading loop in DiffusionPipeline.from_pretrained,
but monkeypatches from_pretrained on ModelMixin and PreTrainedModel to
initialize with random weights instead of loading checkpoint files.
Components with vendored configs use those; others load configs from
the upstream HF model.
"""

import os
import tempfile
from pathlib import Path
from unittest.mock import patch

import torch
from diffusers import ModelMixin
from diffusers.pipelines.pipeline_loading_utils import _get_pipeline_class, simple_get_class_obj
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from transformers import AutoConfig, PreTrainedModel

TINY_MODEL_DIR = os.path.join(tempfile.gettempdir(), "vllm-omni-tiny-models")


def _get_tiny_model_path(name: str) -> str:
    path = os.path.join(TINY_MODEL_DIR, name)
    os.makedirs(path, exist_ok=True)
    return path


def _diffusers_from_config(cls, pretrained_model_name_or_path, **kwargs):
    """Replacement for ModelMixin.from_pretrained that initializes random weights."""
    subfolder = kwargs.get("subfolder")
    load_kwargs = {"subfolder": subfolder} if subfolder else {}
    config = cls.load_config(pretrained_model_name_or_path, **load_kwargs)
    return cls.from_config(config)


def _transformers_from_config(cls, pretrained_model_name_or_path, **kwargs):
    """Replacement for PreTrainedModel.from_pretrained that initializes random weights."""
    subfolder = kwargs.get("subfolder")
    load_kwargs = {"subfolder": subfolder} if subfolder else {}
    config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **load_kwargs)
    return cls(config)


def build_tiny_from_configs(pipeline_name: str, model_id: str, configs_dir: str | Path) -> str:
    """Build a tiny diffusion model from vendored configs with random weights.

    Args:
        pipeline_name: Name of the pipeline (used as output directory name).
        model_id: HuggingFace model ID for loading upstream components
            (tokenizer, scheduler) that don't have vendored configs.
        configs_dir: Path to the directory containing vendored config files
            (model_index.json and per-component config.json).

    Returns:
        Path to the saved tiny model directory with safetensors weights.
    """
    model_dir = _get_tiny_model_path(pipeline_name)
    config_dir = Path(configs_dir)

    config_dict = DiffusionPipeline.load_config(config_dir)
    pipeline_cls = _get_pipeline_class(DiffusionPipeline, config=config_dict)

    init_dict, _, _ = pipeline_cls.extract_init_dict(config_dict)

    # Pop non-component entries (optional pipeline kwargs like is_distilled),
    # same as DiffusionPipeline.from_pretrained lines 345-350
    _, optional_kwargs = DiffusionPipeline._get_signature_keys(pipeline_cls)
    init_kwargs = {k: init_dict.pop(k) for k in optional_kwargs if k in init_dict}

    with (
        patch.object(ModelMixin, "from_pretrained", classmethod(_diffusers_from_config)),
        patch.object(PreTrainedModel, "from_pretrained", classmethod(_transformers_from_config)),
    ):
        for name, (library_name, class_name) in init_dict.items():
            cls = simple_get_class_obj(library_name, class_name)

            # Use vendored config dir if available, otherwise upstream model
            if (config_dir / name).exists():
                init_kwargs[name] = cls.from_pretrained(config_dir / name)
            else:
                init_kwargs[name] = cls.from_pretrained(model_id, subfolder=name)

    pipe = pipeline_cls(**init_kwargs)
    pipe.to(torch.bfloat16).save_pretrained(model_dir)
    return model_dir