import contextlib
from typing import Iterator, Any, Union, Dict
from pathlib import Path
import json

Json = Dict[str, Any]
PathLike = Union[str, Path]


@contextlib.contextmanager
def update_json(path: PathLike) -> Iterator[Json]:
    exists = Path(path).exists()
    with open(path, "r+" if exists else "w") as fo:
        data = json.load(fo) if exists else {}
        yield data
        # Do this in two separate stages so that if there is a JSON decode
        # error, we do not overwrite the file.
        json_data = json.dumps(data, indent=2)
        fo.seek(0)
        fo.write(json_data)
        fo.truncate()


def write_system_metrics(path: PathLike, data: Json) -> None:
    with update_json(f"{path}/metadata.json") as orig_data:
        metrics = orig_data.get("metrics", {})
        system = metrics.get("system", {})
        system.update(data)
        metrics["system"] = system
        orig_data["metrics"] = metrics
