"""Modified from https://github.com/khanrc/honeybee
"""

import os

from peft.utils import ModulesToSaveWrapper
from transformers import logging


def get_cache_dir():
    DEFAULT_HF_HOME = "~/.cache/huggingface"
    cache_dir = os.environ.get("HF_HOME", DEFAULT_HF_HOME)

    return cache_dir


def check_local_file(model_name_or_path):
    cache_dir = get_cache_dir()
    file_name = os.path.join(
        cache_dir, f"models--{model_name_or_path.replace('/', '--')}"
    )
    local_files_only = os.path.exists(file_name)
    file_name = file_name if local_files_only else model_name_or_path
    return local_files_only, file_name


def unwrap_peft(layer):
    """This function is designed for the purpose of checking dtype of model or fetching model configs."""
    if isinstance(layer, ModulesToSaveWrapper):
        return layer.original_module
    else:
        return layer


class transformers_log_level:
    """https://github.com/huggingface/transformers/issues/5421#issuecomment-1317784733
    Temporary set log level for transformers
    """

    orig_log_level: int
    log_level: int

    def __init__(self, log_level: int):
        self.log_level = log_level
        self.orig_log_level = logging.get_verbosity()

    def __enter__(self):
        logging.set_verbosity(self.log_level)

    def __exit__(self, type, value, trace_back):
        logging.set_verbosity(self.orig_log_level)


def get_rank():
    return int(os.environ.get("RANK", 0))