import json
import os
import shutil
import uuid
from contextlib import contextmanager
from pathlib import Path
from typing import Union, Iterable

from ruamel.yaml import YAML


def temp_dir(tmpdir):
    return tmpdir.mkdir(str(uuid.uuid4()))


def temp_path(tmpdir):
    return tmpdir.join(str(uuid.uuid4()))


def temp_json(tmpdir, data):
    return write_json(temp_path(tmpdir), data)


def write_json(path, data):
    with open(path, "w+") as f:
        json.dump(data, f)
    return path


def temp_yaml(tmpdir, data, typ="rt"):
    return write_yaml(temp_path(tmpdir), data, typ)


def write_yaml(path, data, typ="rt"):
    with open(path, "w+") as f:
        YAML(typ=typ).dump(data, f)
    return path


def rmtree_except(path: Path, items_to_preserve: dict):
    if path.exists():
        for item in path.iterdir():
            if item.is_file():
                if item.name not in items_to_preserve:
                    os.unlink(item)
            elif item.is_dir():
                if item.name not in items_to_preserve:
                    shutil.rmtree(item)
                else:
                    sub_except = items_to_preserve[item.name]
                    if isinstance(sub_except, dict):
                        rmtree_except(item, sub_except)


@contextmanager
def change_yaml(path, typ="rt"):
    with open(path, "a+") as f:
        f.seek(0)
        data = YAML(typ=typ).load(f)
        yield data
        f.seek(0)
        f.truncate()
        YAML(typ=typ).dump(data, f)


def read_yaml(path, typ="rt"):
    with open(path, "r") as f:
        return YAML(typ=typ).load(f)


def walk_dir_for_relative_paths(root: Union[str, Path]) -> Iterable[str]:
    if not isinstance(root, Path):
        root = Path(root)
    for parent, _, files in os.walk(root):
        parent = Path(parent)
        for file in files:
            yield str((parent / file).relative_to(root))
