from __future__ import annotations

import json
from collections.abc import Iterable
from dataclasses import dataclass
from typing import Any

import yaml


class DuplicateKeyError(ValueError):
    def __init__(self, key: str, line: int | None = None, column: int | None = None):
        super().__init__(f"Duplicate key {key!r} is not allowed.")
        self.key = key
        self.line = line
        self.column = column


@dataclass(frozen=True, slots=True)
class TaggedValue:
    tag: str
    value: Any


def load_json_document(text: str) -> Any:
    def reject_duplicate_keys(pairs: Iterable[tuple[str, Any]]) -> dict[str, Any]:
        obj: dict[str, Any] = {}
        for key, value in pairs:
            if key in obj:
                raise DuplicateKeyError(key)
            obj[key] = value
        return obj

    return json.loads(text, object_pairs_hook=reject_duplicate_keys)


class CloudFormationLoader(yaml.SafeLoader):
    """PyYAML loader that tolerates CloudFormation intrinsic YAML tags."""


def _freeze_key(value: Any) -> Any:
    if isinstance(value, dict):
        return tuple(sorted((str(key), _freeze_key(item)) for key, item in value.items()))
    if isinstance(value, list):
        return tuple(_freeze_key(item) for item in value)
    return value


def _construct_mapping(loader: CloudFormationLoader, node: yaml.MappingNode, deep: bool = False) -> dict[str, Any]:
    seen: dict[Any, yaml.Node] = {}
    for key_node, _ in node.value:
        key = loader.construct_object(key_node, deep=deep)
        frozen_key = _freeze_key(key)
        if frozen_key in seen:
            raise DuplicateKeyError(str(key), key_node.start_mark.line + 1, key_node.start_mark.column + 1)
        seen[frozen_key] = key_node
    return yaml.SafeLoader.construct_mapping(loader, node, deep=deep)


def _construct_tagged_value(loader: CloudFormationLoader, tag_suffix: str, node: yaml.Node) -> TaggedValue:
    tag = f"!{tag_suffix}"
    if isinstance(node, yaml.ScalarNode):
        return TaggedValue(tag=tag, value=loader.construct_scalar(node))
    if isinstance(node, yaml.SequenceNode):
        return TaggedValue(tag=tag, value=loader.construct_sequence(node))
    if isinstance(node, yaml.MappingNode):
        return TaggedValue(tag=tag, value=loader.construct_mapping(node))
    raise TypeError(f"Unsupported YAML node type for tag {tag}.")


CloudFormationLoader.add_constructor(
    yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
    _construct_mapping,
)
CloudFormationLoader.add_multi_constructor("!", _construct_tagged_value)


def load_yaml_document(text: str) -> Any:
    return yaml.load(text, Loader=CloudFormationLoader)
