"""Lance MLX — first MLX port of ByteDance's unified multimodal model.

Status (2026-05-19):
  modeling_utils      — DONE, parity tested
  vae_wan22           — image mode DONE; video streaming cache PENDING
  lance.py            — backbone + adapter + sampler primitives DONE
  load_lance helper   — loads a converted bundle in one call
"""
from __future__ import annotations

import inspect
import json
from pathlib import Path
from typing import Tuple

import mlx.core as mx

from . import modeling_utils, vae_wan22, lance  # noqa: F401
from .lance import Lance, LanceConfig

__version__ = "0.0.2"


def load_lance(bundle_dir: str | Path) -> Tuple[Lance, dict]:
    """Build Lance from a converted bundle directory and load all weights.

    Expected files in `bundle_dir`:
      config.json
      model.safetensors
      vit.safetensors   (optional — needed for understanding / image-conditioned)
      vae.safetensors   (optional — needed for encode/decode latent <-> pixel)

    Returns (model, cfg_dict).
    """
    from mlx_vlm.models.qwen2_5_vl.config import (
        ModelConfig as Qwen25VLConfig,
        TextConfig,
        VisionConfig,
    )

    bundle_dir = Path(bundle_dir)
    cfg = json.loads((bundle_dir / "config.json").read_text())
    qcfg_dict = cfg["qwen2_5_vl_config"]

    text_sig = inspect.signature(TextConfig).parameters
    text_params = {
        k: v for k, v in qcfg_dict.items()
        if k != "vision_config" and k in text_sig
    }
    text_cfg = TextConfig(**text_params)

    vision_sig = inspect.signature(VisionConfig).parameters
    vision_params = {
        k: v for k, v in qcfg_dict["vision_config"].items()
        if k in vision_sig
    }
    vision_cfg = VisionConfig(**vision_params)

    model_sig = inspect.signature(Qwen25VLConfig).parameters
    model_params = {k: v for k, v in qcfg_dict.items() if k in model_sig}
    model_params["text_config"] = text_cfg
    model_params["vision_config"] = vision_cfg
    qcfg = Qwen25VLConfig(**model_params)

    lance_kwargs = {k: v for k, v in cfg.items() if k != "qwen2_5_vl_config"
                    and k in inspect.signature(LanceConfig).parameters}
    if "latent_patch_size" in lance_kwargs and isinstance(
        lance_kwargs["latent_patch_size"], list
    ):
        lance_kwargs["latent_patch_size"] = tuple(lance_kwargs["latent_patch_size"])
    model = Lance(LanceConfig(qwen_config=qcfg, **lance_kwargs))

    # Bundle key prefixes are Qwen2.5-VL convention (language_model.*, vision_tower.*)
    # plus Lance adapter top-levels (vae2llm, llm2vae, time_embedder, latent_pos_embed)
    # plus vae_model.* for the Wan VAE shards. Our wrapper stores Qwen inside
    # self.backbone, so language_model/vision_tower must be prefixed with `backbone.`.
    # The VAE lives outside the Lance class (separate module), so vae_model.* keys
    # are filtered here and returned via cfg['_vae_weights'] for the caller.
    raw_weights: dict[str, mx.array] = {}
    for fname in ("model.safetensors", "vit.safetensors", "vae.safetensors"):
        path = bundle_dir / fname
        if path.exists():
            raw_weights.update(mx.load(str(path)))

    lance_weights: dict[str, mx.array] = {}
    vae_weights: dict[str, mx.array] = {}
    moe_gen_weights: dict[str, mx.array] = {}
    for k, v in raw_weights.items():
        if k.startswith("vae_model."):
            vae_weights[k[len("vae_model."):]] = v
        elif "_moe_gen" in k:
            # Lance's generation-path weights — not consumed by the bare-Qwen
            # backbone in this wrapper. Returned separately so a future
            # MoE-aware port can pick them up.
            moe_gen_weights[k] = v
        elif k.startswith("language_model.") or k.startswith("vision_tower."):
            lance_weights["backbone." + k] = v
        else:
            lance_weights[k] = v

    model.load_weights(list(lance_weights.items()), strict=False)
    cfg["_vae_weights"] = vae_weights
    cfg["_moe_gen_weights"] = moe_gen_weights
    cfg["_loaded_into_model"] = len(lance_weights)
    return model, cfg