"""Module for wandb utilities"""

import os

from axolotl.utils.dict import DictDefault


def setup_wandb_env_vars(cfg: DictDefault):
    for key in cfg.keys():
        if key.startswith("wandb_"):
            value = cfg.get(key, "")

            if value and isinstance(value, str) and len(value) > 0:
                os.environ[key.upper()] = value

    # Enable wandb if project name is present
    if cfg.wandb_project and len(cfg.wandb_project) > 0:
        cfg.use_wandb = True
        os.environ.pop("WANDB_DISABLED", None)  # Remove if present
    else:
        os.environ["WANDB_DISABLED"] = "true"