## TODO: 增加对于TBStarsParscale的支持
from typing import Dict, List
from collections import OrderedDict

from collators import COLLATORS
from loaders import LOADERS


MODULE_KEYWORDS: Dict[str, Dict[str, List]] = {
    "tbstars_ecom_emb_3b": {
        "vision_encoder": ["model.vision_tower"],
        "vision_projector": ["model.mm_projector"],
        "llm": ["model.model", "model.lm_head"]
    },
    "qwen3vl": {
        "vision_encoder": ["model.visual"],
        "llm": ["model.language_model"]
    },
}


MODEL_HF_PATH = OrderedDict()
MODEL_FAMILIES = OrderedDict()


def register_model(model_id: str, model_family_id: str, model_hf_path: str) -> None:
    if model_id in MODEL_HF_PATH or model_id in MODEL_FAMILIES:
        raise ValueError(f"Duplicate model_id: {model_id}")
    MODEL_HF_PATH[model_id] = model_hf_path
    MODEL_FAMILIES[model_id] = model_family_id


#=============================================================
register_model(
    model_id="tbstars_ecom_emb_3b",
    model_family_id="tbstars_ecom_emb_3b",
    model_hf_path="nebula_internal.TBStars-VL-Ecom-EMB/version=v1"
)

register_model(
    model_id="qwen3vl",
    model_family_id="qwen3vl",
    model_hf_path=""
)


# sanity check
for model_family_id in MODEL_FAMILIES.values():
    if model_family_id != "qwen3vl" and model_family_id not in LOADERS:
        import warnings
        warnings.warn(f"Loader not found for model family: {model_family_id}, it will be unavailable")
    assert model_family_id in MODULE_KEYWORDS, f"Module keywords not found for model family: {model_family_id}"


if __name__ == "__main__":
    temp = "Model ID"
    ljust = 30
    print("Supported models:")
    print(f"  {temp.ljust(ljust)}: HuggingFace Path")
    print("  ------------------------------------------------")
    for model_id, model_hf_path in MODEL_HF_PATH.items():
        print(f"  {model_id.ljust(ljust)}: {model_hf_path}")