import contextlib
import logging
from pathlib import Path

import torch


class _Log:
    def __init__(self):
        logging.basicConfig(level=logging.INFO, format='%(message)s')
        self._logger = logging.getLogger('anydata.vae')

    def info(self, msg):
        self._logger.info(msg)

    def warning(self, msg):
        self._logger.warning(msg)

    def error(self, msg):
        self._logger.error(msg)

    def success(self, msg):
        self._logger.info(msg)


log = _Log()


def get_rank():
    return 0


def broadcast(val, src=0):
    return val


def sync_model_states(*args, **kwargs):
    return None


class _EasyIO:
    @staticmethod
    def load(path, backend_key=None, map_location="cpu"):
        del backend_key
        return torch.load(path, map_location=map_location, weights_only=False)

    @staticmethod
    def save(obj, path, backend_key=None):
        del backend_key
        p = Path(path)
        p.parent.mkdir(parents=True, exist_ok=True)
        torch.save(obj, p)

    @staticmethod
    @contextlib.contextmanager
    def open(path, mode='rb'):
        p = Path(path)
        p.parent.mkdir(parents=True, exist_ok=True)
        with open(p, mode) as f:
            yield f


easy_io = _EasyIO()