"""
Custom YAML document loader.

LangChain doesn't have a dedicated YAML loader, so we provide one here.
This loader reads YAML files and converts them to readable text documents.
"""

from pathlib import Path
from typing import Iterator

import yaml  # type: ignore[import-untyped]
from langchain_core.document_loaders import BaseLoader
from langchain_core.documents import Document
from loguru import logger


class YAMLLoader(BaseLoader):
    """
    Load YAML files and convert to LangChain Documents.

    The YAML content is rendered as formatted YAML text for human readability
    and semantic search. Complex structures are preserved in their YAML format.
    """

    def __init__(self, file_path: str | Path, encoding: str = "utf-8"):
        """
        Initialize the YAML loader.

        Args:
            file_path: Path to the YAML file
            encoding: File encoding (default: utf-8)
        """
        self.file_path = Path(file_path)
        self.encoding = encoding

    def lazy_load(self) -> Iterator[Document]:
        """
        Lazily load documents from the YAML file.

        Yields:
            Document objects containing the YAML content as text
        """
        try:
            with open(self.file_path, encoding=self.encoding) as f:
                content = f.read()

            # Parse YAML to validate and normalize
            try:
                data = yaml.safe_load(content)
            except yaml.YAMLError:
                logger.exception(f"Invalid YAML in {self.file_path}")
                # Still yield the raw content if parsing fails
                yield Document(
                    page_content=content,
                    metadata={
                        "source": str(self.file_path),
                        "parse_error": True,
                    },
                )
                return

            # Skip empty YAML documents
            if data is None:
                logger.debug(f"Skipping empty YAML document: {self.file_path}")
                return

            # Convert back to formatted YAML for readability
            # Use default_flow_style=False for block-style output
            text = yaml.dump(
                data,
                default_flow_style=False,
                allow_unicode=True,
                sort_keys=False,
            )

            yield Document(
                page_content=text,
                metadata={
                    "source": str(self.file_path),
                    "file_type": "yaml",
                },
            )

        except Exception:
            logger.exception(f"Error loading YAML file: {self.file_path}")
            raise

    def load(self) -> list[Document]:
        """Load all documents from the YAML file."""
        return list(self.lazy_load())


def extract_text_from_yaml(content: bytes, encoding: str = "utf-8") -> str:
    """
    Extract text from YAML bytes content.

    This is a convenience function for use with bytes (e.g., from file uploads).

    Args:
        content: YAML content as bytes
        encoding: Text encoding (default: utf-8)

    Returns:
        Formatted YAML text
    """
    try:
        text = content.decode(encoding)
    except UnicodeDecodeError:
        text = content.decode(encoding, errors="ignore")
        logger.warning(
            "YAML content had encoding errors, some characters may be lost"
        )

    try:
        data = yaml.safe_load(text)
        # Skip empty YAML documents
        if data is None:
            return ""
        return yaml.dump(  # type: ignore[no-any-return]
            data,
            default_flow_style=False,
            allow_unicode=True,
            sort_keys=False,
        )
    except yaml.YAMLError:
        logger.exception("Invalid YAML content")
        # Return raw text if parsing fails
        return text
