from typing import TYPE_CHECKING

from pydantic import BaseModel

if TYPE_CHECKING:
    from docling_core.types.doc.document import DocItem, DoclingDocument


class ChunkMetadata(BaseModel):
    """
    Structured metadata for a chunk, including DoclingDocument references.

    Attributes:
        doc_item_refs: JSON pointer references to DocItems in the parent DoclingDocument
                       (e.g., ["#/texts/5", "#/texts/6", "#/tables/0"])
        headings: Section heading hierarchy for this chunk
                  (e.g., ["Chapter 1", "Section 1.1"])
        labels: Semantic labels for each doc_item (e.g., ["paragraph", "table"])
        page_numbers: Page numbers where the chunk content appears
    """

    doc_item_refs: list[str] = []
    headings: list[str] | None = None
    labels: list[str] = []
    page_numbers: list[int] = []

    def resolve_doc_items(self, docling_document: "DoclingDocument") -> list["DocItem"]:
        """Resolve doc_item_refs to actual DocItem objects.

        Args:
            docling_document: The parent DoclingDocument containing the items.

        Returns:
            List of resolved DocItem objects. Items that fail to resolve are skipped.
        """
        from docling_core.types.doc.document import RefItem

        doc_items = []
        for ref in self.doc_item_refs:
            try:
                ref_item = RefItem.model_validate({"$ref": ref})
                doc_item = ref_item.resolve(docling_document)
                doc_items.append(doc_item)
            except Exception:
                # Graceful degradation: skip refs that can't be resolved
                continue
        return doc_items


class Chunk(BaseModel):
    """
    Represents a chunk with content, metadata, and optional document information.
    """

    id: str | None = None
    document_id: str | None = None
    content: str
    metadata: dict = {}
    order: int = 0
    document_uri: str | None = None
    document_title: str | None = None
    document_meta: dict = {}
    embedding: list[float] | None = None

    def get_chunk_metadata(self) -> ChunkMetadata:
        """Parse metadata dict into structured ChunkMetadata."""
        return ChunkMetadata.model_validate(self.metadata)