from .models_config import *

def get_context_tokens():
    context_tokens = {}
    for module_key,values in MODEL_REGISTRY.items():
        context_tokens[module_key] = values.model_max_length
    return context_tokens

DEFAULT_CONTEXT_TOKENS_BY_MODEL: dict[str, int] = get_context_tokens()

def default_context_tokens_for_model(model_key: str) -> int:
    return DEFAULT_CONTEXT_TOKENS_BY_MODEL.get(model_key, 8192)


def get_models_dict_by_tasks(tasks=None):
    tasks = make_list(tasks or [])
    models = {}
    for module_key, values in MODEL_REGISTRY.items():
        # was: if values.task in tasks
        if any(t in tasks for t in values.tasks):
            models[module_key] = values
    return models
def get_models_dict_by_names(names=None):
    names = make_list(names or [])
    models = {}
    for module_key,values in MODEL_REGISTRY.items():
        for name in names:
            if name in values.name:
                models[module_key] = values
                break
    return models


CHAT_MODELS_REGISTRY: Dict[str, ModelConfig] = get_models_dict_by_tasks(tasks=["text-generation","text-generation-inference","text2text-generation"])
if not CHAT_MODELS_REGISTRY.get(DEFAULT_CHAT_MODEL) and CHAT_MODELS_REGISTRY:
    DEFAULT_CHAT_MODEL = list(CHAT_MODELS_REGISTRY.keys())[0]
DEFAULT_MODEL = DEFAULT_CHAT_MODEL


VISION_MODELS_REGISTRY: Dict[str, ModelConfig] = get_models_dict_by_tasks(tasks=["image-text-to-text","text-to-image"])
if not VISION_MODELS_REGISTRY.get(DEFAULT_VISION_MODEL) and VISION_MODELS_REGISTRY:
    DEFAULT_VISION_MODEL = list(VISION_MODELS_REGISTRY.keys())[0]


WHISPER_MODELS_REGISTRY: Dict[str, ModelConfig] = get_models_dict_by_tasks(tasks=["automatic-speech-recognition","speech-recognition"])
if not WHISPER_MODELS_REGISTRY.get(DEFAULT_WHISPER_MODEL) and WHISPER_MODELS_REGISTRY:
    DEFAULT_WHISPER_MODEL = list(WHISPER_MODELS_REGISTRY.keys())[0]


EMBED_MODELS_REGISTRY: Dict[str, ModelConfig] = get_models_dict_by_tasks(tasks=["feature-extraction", "sentence-similarity","sentence-transformers"])
if not EMBED_MODELS_REGISTRY.get(DEFAULT_EMBED_MODEL) and EMBED_MODELS_REGISTRY:
    DEFAULT_EMBED_MODEL = list(EMBED_MODELS_REGISTRY.keys())[0]


# The task-filtered registries above are import-time SNAPSHOTS of MODEL_REGISTRY.
# A model learned at runtime (e.g. pulled/registered from central after startup)
# lands in MODEL_REGISTRY but NOT in these derived dicts — so a vision/whisper/
# embed model resolved after boot raises "Unknown <kind> model key ...
# Available: []". refresh_task_registries() rebuilds them from the current
# MODEL_REGISTRY. It updates IN PLACE (clear+update), which is required because
# other modules import these dicts by reference — rebinding the names here would
# not reach them.
_TASK_REGISTRY_GROUPS = (
    (CHAT_MODELS_REGISTRY,    ["text-generation", "text-generation-inference", "text2text-generation"]),
    (VISION_MODELS_REGISTRY,  ["image-text-to-text", "text-to-image"]),
    (WHISPER_MODELS_REGISTRY, ["automatic-speech-recognition", "speech-recognition"]),
    (EMBED_MODELS_REGISTRY,   ["feature-extraction", "sentence-similarity", "sentence-transformers"]),
)


def refresh_task_registries() -> None:
    """Rebuild the task-filtered registries in place from the live MODEL_REGISTRY.

    Update-then-prune (never clear()): the worker serves on a threaded server, so
    a concurrent reader must never catch the dict momentarily empty. update()
    adds/refreshes current models first, then stale keys are popped one by one —
    so any valid key (notably the one being resolved right now) stays present
    throughout.
    """
    for registry, tasks in _TASK_REGISTRY_GROUPS:
        fresh = get_models_dict_by_tasks(tasks=tasks)
        registry.update(fresh)
        for stale in [k for k in registry if k not in fresh]:
            registry.pop(stale, None)
