#!/usr/bin/env python3
"""Canonical YAML normalization used by every HA import path.

One function, one style, used everywhere data lands in the repo (see
services/home-assistant/DESIGN.md, "Canonical format"): block collections,
indent width 2, unlimited line width, sorted keys, UTF-8.

HA configs use custom `!`-prefixed tags (`!include`, `!secret`, `!env_var`,
`!input`, the `!include_dir_*` family, and any future addition) that
`yaml.safe_load` has no constructor for and rejects with a ConstructorError.
These tags are opaque to us — we never need to resolve `!include`, we just
need to carry it through a parse/dump round-trip unchanged. `TaggedValue`
plus the loader/dumper hooks below handle any `!xxx` tag generically
(scalar, sequence, or mapping payload) without enumerating known tag names.
"""
import json
import sys

import yaml


class TaggedValue:
    """Wraps a value parsed from a custom `!tag` node (see module docstring).

    `tag` is the full tag text (e.g. `"!include"`); `value` is the plain
    Python object (str/list/dict) constructed from the node's payload.
    """

    def __init__(self, tag, value):
        self.tag = tag
        self.value = value

    def __eq__(self, other):
        return (
            isinstance(other, TaggedValue)
            and self.tag == other.tag
            and self.value == other.value
        )

    def __repr__(self):
        return f"TaggedValue({self.tag!r}, {self.value!r})"


class _CanonicalLoader(yaml.SafeLoader):
    pass


def _construct_tagged_value(loader, tag_suffix, node):
    tag = "!" + tag_suffix
    if isinstance(node, yaml.ScalarNode):
        value = loader.construct_scalar(node)
    elif isinstance(node, yaml.SequenceNode):
        value = loader.construct_sequence(node, deep=True)
    elif isinstance(node, yaml.MappingNode):
        value = loader.construct_mapping(node, deep=True)
    else:
        raise yaml.constructor.ConstructorError(
            None, None, f"unsupported node type for tag {tag!r}", node.start_mark
        )
    return TaggedValue(tag, value)


_CanonicalLoader.add_multi_constructor("!", _construct_tagged_value)


class _CanonicalDumper(yaml.SafeDumper):
    pass


def _represent_none(dumper, _value):
    return dumper.represent_scalar("tag:yaml.org,2002:null", "")


def _represent_tagged_value(dumper, data):
    if isinstance(data.value, dict):
        return dumper.represent_mapping(data.tag, data.value)
    if isinstance(data.value, list):
        return dumper.represent_sequence(data.tag, data.value)
    return dumper.represent_scalar(data.tag, data.value)


_CanonicalDumper.add_representer(type(None), _represent_none)
_CanonicalDumper.add_representer(TaggedValue, _represent_tagged_value)


def dump_canonical_yaml(data):
    """Render `data` as canonical YAML text."""
    return yaml.dump(
        data,
        Dumper=_CanonicalDumper,
        default_flow_style=False,
        sort_keys=True,
        allow_unicode=True,
        width=1 << 31,
        indent=2,
    )


def parse_yaml_text(text):
    """Parse YAML text, resolving HA's custom `!include` / `!secret` / etc.

    tags to `TaggedValue` instances instead of raising ConstructorError
    (what plain `yaml.safe_load` does).
    """
    return yaml.load(text, Loader=_CanonicalLoader)


def normalize_yaml_text(text):
    """Parse YAML text then re-dump it through the canonical formatter."""
    data = parse_yaml_text(text)
    return dump_canonical_yaml(data)


def normalize_json_text(text):
    """Parse JSON text (e.g. a `.storage/*` file) and dump as canonical YAML."""
    data = json.loads(text)
    return dump_canonical_yaml(data)


def main(argv):
    if len(argv) != 3 or argv[1] not in ("yaml", "json"):
        print("usage: normalize.py <yaml|json> <input-path>", file=sys.stderr)
        return 2
    mode, path = argv[1], argv[2]
    with open(path, "r", encoding="utf-8") as f:
        text = f.read()
    out = normalize_yaml_text(text) if mode == "yaml" else normalize_json_text(text)
    sys.stdout.write(out)
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))