"""YAML utility functions for the SynthGenAI package."""

import yaml
from loguru import logger


class YamlUtils:
    """Utility class for YAML operations."""

    @staticmethod
    def extract_content(llm_generated_card: str) -> str:
        """
        Extract the content from a YAML string.

        Args:
            llm_generated_card (str): The YAML string generated by LLM.

        Returns:
            str: The extracted content.

        Raises:
            ValueError: If the YAML string does not contain the expected
                sections.
        """
        sections = llm_generated_card.split("---")
        if len(sections) < 3:
            logger.error("Invalid YAML string: missing sections")
            raise ValueError("Invalid YAML string: missing sections")

        return sections[2]

    @staticmethod
    def merge_metadata(hf_generated_card: str, llm_generated_card: str) -> str:
        """
        Merge the metadata from two YAML strings.

        Args:
            hf_generated_card (str): The YAML string generated by Hugging Face.
            llm_generated_card (str): The YAML string generated by LLM.

        Returns:
            str: The merged metadata YAML string.

        Raises:
            ValueError: If the YAML strings do not contain the expected
                sections.
        """
        try:
            hf_generated_card_metadata = yaml.safe_load(
                hf_generated_card.split("---")[1]
            )
            llm_generated_card_metadata = yaml.safe_load(
                llm_generated_card.split("---")[1]
            )
        except IndexError:
            logger.error("Invalid YAML string: missing sections")
            raise ValueError("Invalid YAML string: missing sections")

        merged_metadata = {
            **hf_generated_card_metadata,
            **llm_generated_card_metadata,
        }

        merged_yaml = yaml.dump(merged_metadata, default_flow_style=False)

        merged_output = f"---\n{merged_yaml}---\n"

        return merged_output
