"""Define base classes for chunking."""

import json
from abc import ABC, abstractmethod
from collections.abc import Iterator
from typing import Any, ClassVar

from pydantic import BaseModel
from typing_extensions import deprecated

from docling_core.transforms.serializer.base import BaseDocSerializer
from docling_core.types.doc import DoclingDocument as DLDocument

DFLT_DELIM = "\n"


class BaseMeta(BaseModel):
    """Chunk metadata base class."""

    excluded_embed: ClassVar[list[str]] = []
    excluded_llm: ClassVar[list[str]] = []

    def _excluded_embed_field_names(self) -> set[str]:
        """Return the Python attribute names for all fields listed in `excluded_embed`.

        `excluded_embed` conventionally stores alias keys (the serialised field
        names). `model_dump(exclude=...)` requires Python attribute names, not
        aliases. This helper resolves each entry to its canonical attribute name
        so that exclusion is correct even when a subclass assigns an alias that
        differs from the Python attribute name.

        Returns:
            Python attribute names to exclude from embedding serialization.
        """
        alias_to_name = {(field_info.alias or name): name for name, field_info in type(self).model_fields.items()}
        return {alias_to_name.get(key, key) for key in self.excluded_embed}

    def export_json_dict(self) -> dict[str, Any]:
        """Helper method for exporting non-None keys to JSON mode.

        Returns:
            dict[str, Any]: The exported dictionary.
        """
        return self.model_dump(mode="json", by_alias=True, exclude_none=True)


class BaseChunk(BaseModel):
    """Chunk base class."""

    text: str
    meta: BaseMeta

    def export_json_dict(self) -> dict[str, Any]:
        """Helper method for exporting non-None keys to JSON mode.

        Returns:
            dict[str, Any]: The exported dictionary.
        """
        return self.model_dump(mode="json", by_alias=True, exclude_none=True)


class BaseChunker(BaseModel, ABC):
    """Chunker base class."""

    delim: str = DFLT_DELIM

    @abstractmethod
    def chunk(self, dl_doc: DLDocument, **kwargs: Any) -> Iterator[BaseChunk]:
        """Chunk the provided document.

        Args:
            dl_doc (DLDocument): document to chunk

        Raises:
            NotImplementedError: in this abstract implementation

        Yields:
            Iterator[BaseChunk]: iterator over extracted chunks
        """
        raise NotImplementedError()

    def contextualize(self, chunk: BaseChunk) -> str:
        """Contextualize the given chunk. This implementation is embedding-targeted.

        Args:
            chunk: chunk to serialize

        Returns:
            str: the serialized form of the chunk
        """
        meta = chunk.meta.model_dump(
            mode="json",
            by_alias=True,
            exclude_none=True,
            exclude=chunk.meta._excluded_embed_field_names(),
        )

        items = []
        for k in meta:
            if isinstance(meta[k], list):
                items.append(self.delim.join([d if isinstance(d, str) else json.dumps(d) for d in meta[k]]))
            else:
                items.append(json.dumps(meta[k]))
        items.append(chunk.text)

        return self.delim.join(items)

    @deprecated("Use contextualize() instead.")
    def serialize(self, chunk: BaseChunk) -> str:
        """Contextualize the given chunk. This implementation is embedding-targeted."""
        return self.contextualize(chunk=chunk)


class BaseChunkExpander(BaseModel, ABC):
    """Base chunk expander."""

    @abstractmethod
    def expand(self, chunk: BaseChunk, dl_doc: DLDocument, serializer: BaseDocSerializer) -> BaseChunk:
        """Expand the given chunk.

        Args:
            chunk: The chunk to expand.
            dl_doc: The DoclingDocument containing this chunk.
            serializer: Serializer to convert document content to text.

        Returns:
            BaseChunk: The expanded chunk.
        """
        ...