##############################################################################
# MIT License
#
# Copyright (c) 2025 Advanced Micro Devices, Inc. All Rights Reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

##############################################################################

from pathlib import Path
from typing import Any, Union

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

# --- Round-trip YAML (for writing) ---
RT_YAML = YAML(typ="rt")
RT_YAML.preserve_quotes = True
RT_YAML.width = 4096  # prevent unwanted line wrapping
RT_YAML.indent(mapping=2, sequence=2, offset=0)
RT_YAML.explicit_start = False
RT_YAML.explicit_end = False

# --- Read-only YAML (safe loader) ---
RO_YAML = YAML(typ="safe")
RO_YAML.width = 4096


def load_yaml(
    filepath: Union[str, Path],
    *,
    round_trip: bool = False,
) -> Any:
    path = Path(filepath)
    if not path.exists():
        raise FileNotFoundError(f"YAML file not found: {path}")

    yaml = RT_YAML if round_trip else RO_YAML

    with open(path, "r", encoding="utf-8") as f:
        return yaml.load(f) or CommentedMap()


def save_yaml(data: Any, filepath: Union[str, Path]) -> None:
    path = Path(filepath)
    path.parent.mkdir(parents=True, exist_ok=True)

    with open(path, "w", encoding="utf-8") as f:
        RT_YAML.dump(data, f)


def strip_existing_header(yaml_data) -> None:
    ca = getattr(yaml_data, "ca", None)
    if not ca or not hasattr(ca, "comment") or ca.comment is None:
        return

    original = ca.comment

    cleaned = []

    for block in original:
        if block is None:
            cleaned.append(None)
            continue

        new_block = [token for token in block if "AUTOGENERATED" not in token.value]

        if not new_block:
            cleaned.append(None)
        else:
            cleaned.append(new_block)

    if len(cleaned) < 2:
        cleaned.extend([None] * (2 - len(cleaned)))

    ca.comment = cleaned