"""
Model registry for document ingestion pipeline chunks.

This module defines the metadata schema for chunks created by the document
ingestion pipeline and exposes it as a model registry that can be used by
ChromaSQL queries.
"""

from typing import Mapping, Optional

from pydantic import BaseModel, Field
from idxr.models import ModelSpec


class DocumentChunk(BaseModel):
    """
    Metadata schema for chunks created by the document ingestion pipeline.

    This model describes the metadata fields available on each chunk,
    which can be queried using ChromaSQL.

    Lineage Fields:
        - chunk_no: Sequential chunk number (0-indexed)
        - page_no: Page number for PDF chunks (1-indexed, optional)
        - url: Source URL for the document

    Technical Fields:
        - chunk_offset_tokens: Token offset where this chunk starts
        - token_count: Number of tokens in the chunk
        - byte_count: Number of bytes in the chunk
        - conversation_id: UUID of the conversation that created this chunk
        - ingested_at: ISO timestamp of when the chunk was ingested
        - mime_type: MIME type of the source document
    """

    # Lineage metadata (for tracing chunks back to source)
    chunk_no: int = Field(
        ..., description="Sequential chunk number (0-indexed) within the document"
    )
    page_no: Optional[int] = Field(
        None, description="Page number for PDF chunks (1-indexed, None for non-PDFs)"
    )
    url: str = Field(..., description="Source URL of the document")

    # Technical metadata
    chunk_offset_tokens: int = Field(
        ..., description="Token offset where this chunk starts in the document"
    )
    token_count: int = Field(..., description="Number of tokens in this chunk")
    byte_count: int = Field(..., description="Number of bytes in this chunk")

    # Ingestion context
    conversation_id: str = Field(..., description="UUID of the conversation")
    ingested_at: str = Field(..., description="ISO timestamp of ingestion")
    mime_type: str = Field(..., description="MIME type of the source document")


# Model registry for ChromaSQL queries
# Maps model name -> ModelSpec describing the model's fields and indexing strategy
MODEL_REGISTRY: Mapping[str, ModelSpec] = {
    "DocumentChunk": ModelSpec(
        model=DocumentChunk,
        semantic_fields=[],  # No semantic fields - chunks are already embedded
        keyword_fields=[
            # Lineage fields for filtering and grouping
            "chunk_no",
            "page_no",
            "url",
            # Technical fields for debugging and analysis
            "chunk_offset_tokens",
            "token_count",
            "byte_count",
            "conversation_id",
            "ingested_at",
            "mime_type",
        ],
    ),
}


def get_model_registry() -> Mapping[str, ModelSpec]:
    """
    Get the model registry for document ingestion pipeline chunks.

    Returns:
        Dictionary mapping model names to ModelSpec instances
    """
    return MODEL_REGISTRY
