"""First-class Pydantic provenance stamp.

Provenance is authoritative framework data — **independent** of MAF's ``Annotation`` type,
which silently drops on the Python streaming path (#4316, research 02 Dim 2). A
``ProvenanceStamp`` is a Pydantic model that must carry at least one ``Citation`` (a
``min_length=1`` constraint), the model + role that produced the proposal, the validator's
decision, and the token usage. ``to_annotations()`` maps to MAF ``Annotation`` dicts for
DISPLAY only — never the source of truth.

The locator type (``TextSpan``) is owned by ``retrieval.py`` (Step 5) and imported here, so
a citation's span is the same exact object the retriever produced.
"""

from __future__ import annotations

from typing import Literal

from agent_framework import Annotation, TextSpanRegion
from pydantic import BaseModel, Field

from portfolio_optimiser.retrieval import TextSpan


class Citation(BaseModel):
    """One exact-span citation into a source document (locator owned by retrieval.py)."""

    file: str
    locator: TextSpan
    snippet: str


class ProvenanceStamp(BaseModel):
    """Authoritative provenance for one proposal — at least one citation is mandatory."""

    citations: list[Citation] = Field(min_length=1)
    model: str
    role: str
    validator_decision: Literal["validated", "rejected"]
    token_usage: int

    def to_annotations(self) -> list[Annotation]:
        """Map to MAF ``Annotation`` dicts for display only (NOT the source of truth)."""
        return [
            Annotation(
                type="citation",
                snippet=c.snippet,
                annotated_regions=[
                    TextSpanRegion(
                        type="text_span",
                        start_index=c.locator.start_index,
                        end_index=c.locator.end_index,
                    )
                ],
                additional_properties={"file": c.file},
            )
            for c in self.citations
        ]