import os
from typing import Literal, overload

from ruamel.yaml import YAML
from ruamel.yaml.comments import CommentedMap


def _create_YAML(preserve_comments: bool = False) -> YAML:
    """
    Creates a YAML that can load/save with comments if preserve_comments is True.
    """
    typ = None if preserve_comments else "safe"
    result = YAML(typ=typ)
    result.preserve_quotes = preserve_comments
    return result


@overload
def load_yaml(path: str, preserve_comments: Literal[False]) -> dict: ...
@overload
def load_yaml(path: str, preserve_comments: Literal[True]) -> CommentedMap: ...
def load_yaml(path: str, preserve_comments: bool = False) -> dict | CommentedMap:
    with open(path, encoding="utf-8") as f:
        yaml = _create_YAML(preserve_comments)
        return yaml.load(f)


def save_yaml(path: str, data: dict | CommentedMap, preserve_comments: bool = False) -> None:
    yaml = _create_YAML(preserve_comments)
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        yaml.dump(data, f)
