from .models_dict import MODELS
from .imports import *
 
# ---------------------------------------------------------------------
# Model registry
# ---------------------------------------------------------------------
 
# Disk overlay may override only these ops knobs; everything else on disk is
# ignored (code is the source of truth for identity/routing fields).
 
 
def assess_config(cls, values):
    """Build cls from values if it can be a valid instance, else False.
 
    Required = fields the dataclass defines with no default. Unknown keys are
    dropped, absent optionals are filled from their defaults. Never raises on
    missing/extra fields.
    """
    flds = {f.name: f for f in fields(cls)}
 
    for f in flds.values():
        required = f.default is MISSING and f.default_factory is MISSING
        if required and values.get(f.name) in (None, "", []):
            return False
 
    out = {}
    for name, f in flds.items():
        if name in values:
            out[name] = values[name]
        elif f.default is not MISSING:
            out[name] = f.default
        else:
            out[name] = f.default_factory()
    return cls(**out)
 
 
def get_model_values(config, dict_return=False, return_dict=False):
    """Return the ModelConfig, or its dict form if dict_return/return_dict."""
    if dict_return or return_dict:
        return config.to_dict()
    return config
 
 
def get_assessed_model_config(values, dict_return=False, return_dict=False):
    assessed = assess_config(ModelConfig, values)
    if assessed is False:
        return False
    return get_model_values(assessed, dict_return=dict_return, return_dict=return_dict)
 
 
def update_model_config_dict(model_key=None, values=None, dict_obj=None,
                             dict_return=False, return_dict=False, key=None):
    dict_obj = dict_obj if dict_obj is not None else {}
    model_key = model_key or key
    values = dict(values or {})          # copy — don't mutate caller's dict
    values["model_key"] = model_key
 
    config = get_assessed_model_config(values, dict_return=dict_return, return_dict=return_dict)
    if config is False:                  # incomplete -> not attached
        return dict_obj
 
    dict_obj[model_key] = config
    return dict_obj
 
 
def get_models_raw_dict(models_dict_path=None):
    models_dict_path = models_dict_path or MODELS_DICT_PATH
    if models_dict_path and os.path.isfile(models_dict_path):
        return safe_load_from_json(models_dict_path) or {}
    return {}
 
 
def get_models_dict(models_dict_path=None, dict_return=False, return_dict=False):
    dict_return = dict_return or return_dict
    overlay = get_models_raw_dict(models_dict_path=models_dict_path)
    nudict = {}
 
    # code-defined models, with whitelisted overlay knobs merged in
    for model_key, defaults in MODELS.items():
        merged = dict(defaults)
        if model_key in overlay:
            merged.update({k: v for k, v in overlay[model_key].items()
                           if k in OVERLAY_ALLOWED})
        nudict = update_model_config_dict(
            model_key=model_key, values=merged, dict_obj=nudict, dict_return=dict_return
        )
 
    # disk-only models (present on disk, absent from code)
    for model_key, values in overlay.items():
        if model_key in MODELS:
            continue
        nudict = update_model_config_dict(
            model_key=model_key, values=values, dict_obj=nudict, dict_return=dict_return
        )
 
    return nudict
 

MODEL_REGISTRY: Dict[str, ModelConfig] = get_models_dict()
MODEL_REGISTRY_DICT: Dict[str, dict] = get_models_dict(dict_return=True)
    
def get_model_registry(dict_return=False,return_dict=False) -> ModelConfig or dict:
    dict_return = dict_return or return_dict
    model_registry = MODEL_REGISTRY
    if dict_return:
        model_registry = MODEL_REGISTRY_DICT
    return model_registry
