"""Load Fibo 2.0 with OrbitQuant W4A8 components into diffusers' ``BriaFibo2Pipeline``."""

from __future__ import annotations

import dataclasses
import json
import os
from pathlib import Path

import torch
from torch import nn

GROUPS = ("qkv", "attn_out", "ff_in", "ff_out", "xq", "xkv", "xout")
# The deepest Qwen3-VL hidden state the pipeline reads is 35, the output of layer 34.
TEXT_LAYERS = 35

# Negative captions of the original model card: text-to-image steers away from soft, featureless, artifact-laden
# images; edits use the empty caption the model was trained with.
NEGATIVE_T2I = '{"short_description":"","objects":[{"description":"the main subject","texture":"smooth, featureless, plastic-like, airbrushed, no fine surface detail","appearance_details":"soft blurred edges, simplified shapes, no fine detail","skin_tone_and_texture":"waxy airbrushed skin, poreless, plastic smooth"}],"background_setting":"a plain seamless studio backdrop, perfectly smooth gradient, blurred and featureless, no visible texture","lighting":{},"aesthetics":{"artifacts":["oversaturated colours","neon colour cast","overexposed highlights","lens flare","slight blur","motion blur","haze"]},"photographic_characteristics":{"depth_of_field":"shallow, heavy bokeh obscuring most of the frame","focus":"soft, slightly out of focus, lacking fine texture"},"style_medium":"","context":"","artistic_style":""}'
NEGATIVE_EDIT = '{"short_description":"","objects":[],"background_setting":"","lighting":{},"aesthetics":{},"photographic_characteristics":{},"style_medium":"","context":"","artistic_style":""}'


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, device: str, fused: bool, attention: str, activation_bits: int):
    import orbitquant  # noqa: F401 -- registers the OrbitQuant layers
    from diffusers import BriaFibo2Transformer2DModel as Transformer
    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.from_config(Transformer.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))
    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

        fuse(model, config=config, attention=attention, activation_bits=dict.fromkeys(GROUPS, activation_bits))
    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 152k rows, so the 0.8 GB 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):
    """Qwen3-VL without the LM head and the unused last decoder layer. The final norm becomes the identity, so the
    last recorded hidden state is the raw output of layer 34, as in the full model; the vision tower stays BF16 for
    the images to edit and the drawn caption texts."""
    import orbitquant  # noqa: F401 -- registers the OrbitQuant backend with Transformers
    from transformers import Qwen3VLForConditionalGeneration

    full = Qwen3VLForConditionalGeneration.from_pretrained(str(path), dtype=torch.bfloat16, device_map={"": device})
    model = full.model.eval()
    language = model.language_model
    language.layers = language.layers[:TEXT_LAYERS]
    language.config.num_hidden_layers = TEXT_LAYERS
    if getattr(language.config, "layer_types", None):
        language.config.layer_types = language.config.layer_types[:TEXT_LAYERS]
    language.norm = nn.Identity()
    del full
    if fused:
        from orbitquant.fused import fuse_linears

        fuse_linears(model, activation_bits=activation_bits)
    if host_embedding:
        language.embed_tokens = HostEmbedding(language.embed_tokens)
    torch.cuda.empty_cache()
    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):
    """Fibo 2.0 pipeline with the W4A8 transformer and text encoder; call it as the original ``BriaFibo2Pipeline``.

    ``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)
    from diffusers import AutoencoderKLFlux2, BriaFibo2Pipeline, FlowMatchEulerDiscreteScheduler
    from transformers import Qwen3VLProcessor

    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", 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"), dtype=torch.bfloat16)
    vae = vae.to("cpu" if offload else device).eval()
    pipeline = BriaFibo2Pipeline(
        transformer=transformer,
        scheduler=FlowMatchEulerDiscreteScheduler.from_pretrained(str(path / "scheduler")),
        vae=vae,
        text_encoder=text_encoder,
        processor=Qwen3VLProcessor.from_pretrained(str(path / "processor")),
    )
    if offload:
        pipeline.enable_model_cpu_offload(device=device)
    torch.cuda.empty_cache()
    return pipeline


def fonts(path: str | Path) -> list[str]:
    """The fonts Fibo 2.0's training drew caption texts with (pass them as ``fonts=``)."""
    return [str(Path(path) / "fonts" / name) for name in ("Aileron-Regular.otf", "DejaVuSans.ttf")]


def sampling(path: str | Path) -> dict:
    """Default steps and guidance of the release (Dev: 50 and 5.0; Edit Turbo: 4 and 1.0)."""
    return json.loads((Path(path) / "manifest.json").read_text())["sampling"]