"""fusion-embedding-2-k3-vision — Kimi K3's vision encoder, made language-searchable.

A small trained projector maps Kimi K3's frozen MoonViT-V2 vision features into the
frozen fusion-embedding-2 (Qwen3-VL-Embedding-2B) text space, so K3's own visual
features become searchable in language and sit in the same space as fusion-embedding's
text, image, video, and audio. This ships the projector only; the ~401M vision tower is
pulled from moonshotai/Kimi-K3 on first use (never the 2.8T language model).

    from inference import K3VisionEmbedder
    m = K3VisionEmbedder.from_pretrained("EximiusLabs/fusion-embedding-2-k3-vision")
    v = m.embed_image("photo.jpg")                       # 2048-d, in the FE2 text space

Query with text using fusion-embedding-2 itself:
    from inference import FusionEmbedder            # from the fusion-embedding-2 repo
    fe = FusionEmbedder.from_pretrained("EximiusLabs/fusion-embedding-2-2b-preview")
    q = fe.embed_text("a dog on a beach")
    score = (q @ v)                                       # cosine, both unit-norm

MoonViT-V2 is not published as a standalone checkpoint; it lives inside Kimi-K3 with a
`vision_tower.` prefix. We load only those tensors from the shard index. Loading the
vision tower's modeling code also imports K3's language stack, so we stub that import
(the vision path never uses it) rather than pulling in the LLM's kernels.

Requires: torch, transformers>=4.57, einops, pillow, safetensors, huggingface_hub.
The Kimi-K3 base is under the Kimi K3 License (permissive, notice required); this
projector is released under CC-BY-NC-4.0.
"""

from __future__ import annotations

import os
import sys
import json
import types
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

K3_REPO = "moonshotai/Kimi-K3"
BASE = "Qwen/Qwen3-VL-Embedding-2B"
VPREFIX = "vision_tower."


# --------------------------------------------------------------------------- #
# MoonViT-V2 loader: only the vision tower, no language model, no fla/triton
# --------------------------------------------------------------------------- #
def _install_llm_stub() -> None:
    """Stub K3's language modeling submodule so importing the ViT never pulls the LLM."""
    from importlib.abc import Loader, MetaPathFinder
    from importlib.machinery import ModuleSpec

    class _StubLoader(Loader):
        def create_module(self, spec):
            m = types.ModuleType(spec.name)
            m.__getattr__ = lambda name: type(name, (), {})
            return m

        def exec_module(self, module):
            pass

    class _StubFinder(MetaPathFinder):
        def find_spec(self, fullname, path=None, target=None):
            if fullname.endswith("modeling_kimi_linear"):
                return ModuleSpec(fullname, _StubLoader())
            return None

    if not any(type(f).__name__ == "_StubFinder" for f in sys.meta_path):
        sys.meta_path.insert(0, _StubFinder())


def load_moonvit_tower(device: str, dtype=torch.bfloat16, token: Optional[str] = None):
    """Return the frozen MoonViT-V2 tower and its image processor."""
    import safetensors.torch as st
    from huggingface_hub import hf_hub_download
    from transformers import AutoConfig, AutoImageProcessor
    from transformers.dynamic_module_utils import get_class_from_dynamic_module

    token = token or os.environ.get("HF_TOKEN")
    cfg = AutoConfig.from_pretrained(K3_REPO, trust_remote_code=True, token=token)
    vcfg = cfg.vision_config

    # the LLM stack imports transformers-internal symbols this build may not expose;
    # shim them (only referenced by the LLM, which the ViT never runs).
    import transformers.utils.generic as _g
    if not hasattr(_g, "OutputRecorder"):
        _g.OutputRecorder = type("OutputRecorder", (), {"__init__": lambda self, *a, **k: None})
    if not hasattr(_g, "check_model_inputs"):
        _g.check_model_inputs = lambda func=None, *a, **k: (func if callable(func) else (lambda f: f))

    _install_llm_stub()
    MoonViT = get_class_from_dynamic_module(
        "modeling_kimi_k3.MoonViT3dPretrainedModel", K3_REPO, token=token)

    # KimiK3VisionConfig prefixes vision fields with vt_; alias to the standard names
    # (hidden_size is the ViT width 1024, not the 7168 LLM width).
    for _std, _vt in (("hidden_size", "vt_hidden_size"),
                      ("num_hidden_layers", "vt_num_hidden_layers"),
                      ("num_attention_heads", "vt_num_attention_heads"),
                      ("intermediate_size", "vt_intermediate_size")):
        if not hasattr(vcfg, _std) and hasattr(vcfg, _vt):
            setattr(vcfg, _std, getattr(vcfg, _vt))
    vcfg._attn_implementation = "eager"      # the ViT attn table has no sdpa path
    tower = MoonViT(vcfg)

    # load only vision_tower.* tensors from the shard index (never the MoE experts)
    idx = json.load(open(hf_hub_download(K3_REPO, "model.safetensors.index.json", token=token)))
    wmap = idx["weight_map"]
    vkeys = {k: v for k, v in wmap.items() if k.startswith(VPREFIX)}
    sd = {}
    for sh in sorted(set(vkeys.values())):
        t = st.load_file(hf_hub_download(K3_REPO, sh, token=token))
        for k in vkeys:
            if wmap[k] == sh:
                sd[k[len(VPREFIX):]] = t[k]
    miss, unexp = tower.load_state_dict(sd, strict=False)
    if miss or unexp:
        raise RuntimeError(f"MoonViT load mismatch: missing={len(miss)} unexpected={len(unexp)}")

    for m in tower.modules():
        if hasattr(m, "attn_implementation"):
            m.attn_implementation = "eager"
    tower = tower.to(device, dtype).eval()
    for p in tower.parameters():
        p.requires_grad_(False)

    proc = AutoImageProcessor.from_pretrained(K3_REPO, trust_remote_code=True, token=token)
    return tower, proc


def _build_projector(in_dim: int, out_dim: int) -> nn.Sequential:
    return nn.Sequential(nn.LayerNorm(in_dim), nn.Linear(in_dim, 1024), nn.GELU(),
                         nn.Dropout(0.2), nn.Linear(1024, out_dim))


class K3VisionEmbedder:
    """Frozen MoonViT-V2 + a trained projector into the fusion-embedding-2 text space."""

    def __init__(self, ckpt_path: str, device: str = "cuda", dtype=torch.bfloat16):
        ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
        self.cfg = ck["config"]
        self.device = device
        self.tower, self.proc = load_moonvit_tower(device, dtype)
        self.proj = _build_projector(self.cfg["in_dim"], self.cfg["out_dim"]).to(device).float()
        self.proj.load_state_dict({k: v.float() for k, v in ck["proj"].items()})
        self.proj.eval()

    @classmethod
    def from_pretrained(cls, repo_or_path: str, device: str = "cuda",
                        revision: Optional[str] = None, **kw) -> "K3VisionEmbedder":
        if os.path.exists(repo_or_path):
            path = repo_or_path
        else:
            from huggingface_hub import hf_hub_download
            path = hf_hub_download(repo_or_path, "k3vit2fe2_proj.pt", revision=revision)
        return cls(path, device=device, **kw)

    @torch.no_grad()
    def _tower_feat(self, image) -> torch.Tensor:
        from PIL import Image
        if isinstance(image, (str, os.PathLike)):
            image = Image.open(str(image))
        image = image.convert("RGB")
        enc = self.proc.preprocess([{"type": "image", "image": image}], return_tensors="pt")
        pv = enc["pixel_values"].to(self.device, torch.bfloat16)
        grid = enc["grid_thws"].to(self.device)
        f = self.tower(pv, grid)[0].float()               # [N_merged, 4, 1024]
        return F.normalize(f.mean(dim=(0, 1)), dim=-1)     # [1024]

    @torch.no_grad()
    def embed_image(self, image) -> torch.Tensor:
        """Embed an image via MoonViT-V2 into the fusion-embedding-2 text space (2048-d)."""
        feat = self._tower_feat(image)
        return F.normalize(self.proj(feat), dim=-1).cpu()