"""Load MOSS-Image-SFT with OrbitQuant W4A8 components into the model's own pipeline (``pipeline_moss_image.py``)."""

from __future__ import annotations

import dataclasses
import json
import os
import sys
from pathlib import Path

import torch
from torch import nn

GROUPS = ("qkv", "mlp", "attn_out", "mlp_out")


def _allocator(device: str) -> None:
    # Growable segments return the transient blocks of one stage to the next. An explicit setting is left untouched.
    if not device.startswith("cuda") or {"PYTORCH_CUDA_ALLOC_CONF", "PYTORCH_ALLOC_CONF"} & set(os.environ):
        return
    if not torch.cuda.is_initialized():
        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
    else:
        setter = getattr(torch._C, "_accelerator_setAllocatorSettings", None)
        (setter or torch.cuda.memory._set_allocator_settings)("expandable_segments:True")


def _load_transformer(path: Path, transformer_class, device: str, fused: bool, attention: str, activation_bits: int):
    import orbitquant  # noqa: F401 -- registers the OrbitQuant layers
    from orbitquant.adaln import RTNInt4Linear
    from orbitquant.artifacts.manifest import OrbitQuantManifest
    from orbitquant.config import OrbitQuantConfig
    from orbitquant.layers import OrbitQuantLinear
    from orbitquant.modeling import _parent_and_child, _set_child
    from safetensors.torch import load_file

    config = OrbitQuantConfig.from_dict(json.loads((path / "quantization_config.json").read_text()))
    manifest = OrbitQuantManifest.from_dict(json.loads((path / "orbitquant_manifest.json").read_text()))
    with torch.device("meta"):
        model = transformer_class.from_config(transformer_class.load_config(str(path)))
    for name in manifest.quantized_modules:
        bits = manifest.module_bits.get(name)
        layer_config = dataclasses.replace(config, weight_bits=bits) if bits and bits != config.weight_bits else config
        parent, child = _parent_and_child(model, name)
        _set_child(parent, child, OrbitQuantLinear.empty_from_linear(
            model.get_submodule(name), config=layer_config, module_name=name))
    for name in manifest.adaln_modules:
        parent, child = _parent_and_child(model, name)
        _set_child(parent, child, RTNInt4Linear.empty_from_linear(model.get_submodule(name), config=config, module_name=name))
    model.load_state_dict(load_file(str(path / "model.safetensors"), device=device), strict=True, assign=True)
    model = model.to(device).eval()
    if fused:
        from orbitquant.fused import fuse

        widths = dict.fromkeys(GROUPS, activation_bits) if activation_bits != config.activation_bits else None
        fuse(model, config=config, attention=attention, activation_bits=widths)
    return model


def _park(model: nn.Module) -> None:
    """Move a component to host memory right after its GPU build, so the next build does not stack on it."""
    from orbitquant.fused import release

    release(model)
    model.to("cpu")
    torch.cuda.empty_cache()


class HostEmbedding(nn.Module):
    """Token embedding kept in host memory: a prompt reads a few hundred of its 248k rows, so the 2 GiB table stays
    off the GPU. The table is not a registered parameter, so moving the text encoder does not move it."""

    def __init__(self, embedding: nn.Embedding):
        super().__init__()
        self.__dict__["table"] = embedding.weight.detach().to("cpu")
        self.num_embeddings, self.embedding_dim = self.table.shape

    @property
    def weight(self):
        return self.table

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        return nn.functional.embedding(ids.cpu(), self.table).to(ids.device, non_blocking=True)


def _load_text_encoder(path: Path, device: str, fused: bool, activation_bits: int, host_embedding: bool):
    import orbitquant  # noqa: F401 -- registers the OrbitQuant backend with Transformers
    from transformers import AutoModel

    model = AutoModel.from_pretrained(str(path), dtype=torch.bfloat16, device_map={"": device}).eval()
    if fused:
        from orbitquant.fused import fuse_linears

        fuse_linears(model, activation_bits=activation_bits)
    if host_embedding:
        model.embed_tokens = HostEmbedding(model.embed_tokens)
    return model


def load_pipeline(path: str | Path, device: str = "cuda", attention: str = "cudnn", offload: bool = False,
                  fused: bool | None = None, activation_bits: int = 8, host_embedding: bool = True):
    """MOSS-Image pipeline with the W4A8 transformer and text encoder; call it as the original pipeline.

    ``attention`` is ``cudnn`` (default), ``flash``, ``sdpa`` or ``int8`` (faster, changes the image slightly).
    ``offload`` keeps the components in host memory and moves each to the GPU only while it runs. ``fused=None``
    fuses the projections on CUDA devices with compute capability 8.0 or newer (needed for 8-bit activations)."""
    path = Path(path)
    _allocator(device)
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))
    from diffusers import AutoencoderKLFlux2
    from pipeline_moss_image import MOSSImagePipeline, MOSSImageScheduler, MOSSImageTransformer
    from transformers import AutoTokenizer

    if fused is None:
        fused = device.startswith("cuda") and torch.cuda.get_device_capability(device) >= (8, 0)
    # The transformer is fused on the GPU even with offload: on the host the group build runs many small parallel
    # copies, which stall for minutes when a container's CPU quota is far below its visible core count.
    transformer = _load_transformer(path / "transformer", MOSSImageTransformer, device, fused, attention, activation_bits)
    if offload:
        _park(transformer)
    text_encoder = _load_text_encoder(path / "text_encoder", device, fused, activation_bits, host_embedding)
    if offload:
        _park(text_encoder)
    vae = AutoencoderKLFlux2.from_pretrained(str(path / "vae"), torch_dtype=torch.float32)
    vae = vae.to("cpu" if offload else device).eval()
    tokenizer = AutoTokenizer.from_pretrained(str(path / "tokenizer"))
    pipeline = MOSSImagePipeline(transformer, MOSSImageScheduler(), text_encoder, tokenizer, vae)
    if offload:
        pipeline.enable_model_cpu_offload(device=device)
    torch.cuda.empty_cache()
    return pipeline