"""Pinned Sakura Qwen3-VL encoder with the official FLUX.3 processor/config.

Only the shared text/vision encoder changes. The FLUX image transformers and
VAE remain official. Set FLUX3_TEXT_ENCODER=official to restore that encoder.
"""
from __future__ import annotations

import hashlib
import json
import os
from pathlib import Path
import tempfile

OFFICIAL_REPO = "black-forest-labs/flux-3-image-klein-base"
OFFICIAL_REVISION = "3f7fa5ee22ca1d5a5cc095ac7cc96e6a173b4c86"
SAKURA_REPO = "webmp3/Sakura-FLUX.3-Image-Klein-TextEncoder-Uncensored-GGUF"
SAKURA_REVISION = "fe49271a47b61efe0d96b476306c25934d943472"
SAKURA_FILE = "Sakura-TextEncoder-Qwen3VL-4B-Uncensored-bf16-8.27GiB.safetensors"
SAKURA_BYTES = 8_875_719_408
# All 713 names, shapes and dtypes independently match the official encoder.
ENCODER_SCHEMA_SHA256 = "f4619891efac96ad653344f1b195a51427b5cc9637769f58f435d20af26c93dc"
ENCODER_MODE = os.environ.get("FLUX3_TEXT_ENCODER", "sakura").strip().lower()
if ENCODER_MODE not in {"sakura", "official"}:
    raise ValueError("FLUX3_TEXT_ENCODER must be 'sakura' or 'official'.")
ENCODER_LABEL = "Sakura text encoder" if ENCODER_MODE == "sakura" else "Official text encoder"


def text_encoder_info():
    return {
        "variant": ENCODER_MODE,
        "label": ENCODER_LABEL,
        "repo": SAKURA_REPO if ENCODER_MODE == "sakura" else OFFICIAL_REPO,
        "revision": SAKURA_REVISION if ENCODER_MODE == "sakura" else OFFICIAL_REVISION,
        "processor_repo": OFFICIAL_REPO,
        "processor_revision": OFFICIAL_REVISION,
    }


def _validate_checkpoint(checkpoint):
    """Read only the safetensors header; never allocate the 8.27 GiB tensors."""
    path = Path(checkpoint)
    if path.stat().st_size != SAKURA_BYTES:
        raise RuntimeError("The Sakura text encoder checkpoint has an unexpected size.")
    with path.open("rb") as stream:
        header_size = int.from_bytes(stream.read(8), "little")
        if not 2 <= header_size <= 1_000_000:
            raise RuntimeError("The Sakura text encoder checkpoint has an invalid header.")
        try:
            header = json.loads(stream.read(header_size))
            schema = {name: {"dtype": value["dtype"], "shape": value["shape"]}
                      for name, value in header.items() if name != "__metadata__"}
        except (ValueError, TypeError, KeyError, AttributeError) as exc:
            raise RuntimeError("The Sakura text encoder checkpoint header is unreadable.") from exc
    digest = hashlib.sha256(json.dumps(schema, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
    if digest != ENCODER_SCHEMA_SHA256:
        raise RuntimeError("The Sakura text encoder tensor layout differs from the verified Qwen3-VL encoder.")


def _check_loading_info(info, config):
    # Transformers can report an omitted output head even when it is correctly
    # tied to the input embeddings. The checkpoint schema is checked separately.
    allowed_missing = {"lm_head.weight"} if getattr(config, "tie_word_embeddings", False) else set()
    missing = set(info.get("missing_keys", [])) - allowed_missing
    if missing or any(info.get(key) for key in ("unexpected_keys", "mismatched_keys", "error_msgs")):
        raise RuntimeError("Text encoder loading was incomplete; refusing to run with missing or mismatched weights.")


def load_text_encoder(dtype, token=None):
    from huggingface_hub import hf_hub_download
    from transformers import AutoConfig, Qwen3VLForConditionalGeneration

    config = AutoConfig.from_pretrained(
        OFFICIAL_REPO, subfolder="text_encoder", revision=OFFICIAL_REVISION, token=token,
    )
    if ENCODER_MODE == "official":
        model, info = Qwen3VLForConditionalGeneration.from_pretrained(
            OFFICIAL_REPO, subfolder="text_encoder", revision=OFFICIAL_REVISION,
            config=config, dtype=dtype, token=token, output_loading_info=True,
        )
    else:
        checkpoint = hf_hub_download(
            SAKURA_REPO, SAKURA_FILE, revision=SAKURA_REVISION, token=token,
        )
        _validate_checkpoint(checkpoint)
        # Transformers expects model.safetensors beside config.json. A temporary
        # symlink keeps the cached weight file in place without an 8 GiB copy.
        # The Space runs Linux; only this tiny directory is temporary.
        with tempfile.TemporaryDirectory(prefix="flux3-sakura-encoder-") as directory:
            config.save_pretrained(directory)
            os.symlink(str(Path(checkpoint).resolve()), str(Path(directory) / "model.safetensors"))
            model, info = Qwen3VLForConditionalGeneration.from_pretrained(
                directory, config=config, dtype=dtype, local_files_only=True,
                output_loading_info=True,
            )
    _check_loading_info(info, config)
    model.tie_weights()
    return model.eval()