from typing import Dict, List
from collections import OrderedDict

from collators import COLLATORS
from loaders import LOADERS


MODULE_KEYWORDS: Dict[str, Dict[str, List]] = {
    "qwen3-vl": {
        "vision_encoder": ["visual.patch_embed", "visual.rotary_pos_emb", "visual.blocks"],
        "vision_projector": ["visual.merger"],
        "llm": ["language_model"]
    },
    "qwen3_5": {
        "vision_encoder": ["model.visual.patch_embed", "model.visual.pos_embed", "model.visual.blocks"],
        "vision_projector": ["model.visual.merger"],
        "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="qwen3-vl-2b",
    model_family_id="qwen3-vl",
    model_hf_path="Qwen/Qwen3-VL-2B-Instruct"
)
register_model(
    model_id="qwen3-vl-8b",
    model_family_id="qwen3-vl",
    model_hf_path="Qwen/Qwen3-VL-8B-Instruct"
)
register_model(
    model_id="qwen3_5-9b",
    model_family_id="qwen3_5",
    model_hf_path="Qwen/Qwen3.5-VL-9B"
)

for model_family_id in MODEL_FAMILIES.values():
    assert model_family_id in COLLATORS, f"Collator not found for model family: {model_family_id}"
    assert model_family_id in LOADERS, f"Loader not found for model family: {model_family_id}"
    assert model_family_id in MODULE_KEYWORDS, f"Module keywords not found for model family: {model_family_id}"