"""Citation type definitions for the SDK.

These types are modeled after the Bedrock API.
"""

from typing import Literal

from typing_extensions import TypedDict


class CitationsConfig(TypedDict):
    """Configuration for enabling citations on documents.

    Attributes:
        enabled: Whether citations are enabled for this document.
    """

    enabled: bool


class DocumentCharLocation(TypedDict, total=False):
    """Specifies a character-level location within a document.

    Provides precise positioning information for cited content using
    start and end character indices.

    Attributes:
        documentIndex: The index of the document within the array of documents
            provided in the request. Minimum value of 0.
        start: The starting character position of the cited content within
            the document. Minimum value of 0.
        end: The ending character position of the cited content within
            the document. Minimum value of 0.
    """

    documentIndex: int
    start: int
    end: int


class DocumentChunkLocation(TypedDict, total=False):
    """Specifies a chunk-level location within a document.

    Provides positioning information for cited content using logical
    document segments or chunks.

    Attributes:
        documentIndex: The index of the document within the array of documents
            provided in the request. Minimum value of 0.
        start: The starting chunk identifier or index of the cited content
            within the document. Minimum value of 0.
        end: The ending chunk identifier or index of the cited content
            within the document. Minimum value of 0.
    """

    documentIndex: int
    start: int
    end: int


class DocumentPageLocation(TypedDict, total=False):
    """Specifies a page-level location within a document.

    Provides positioning information for cited content using page numbers.

    Attributes:
        documentIndex: The index of the document within the array of documents
            provided in the request. Minimum value of 0.
        start: The starting page number of the cited content within
            the document. Minimum value of 0.
        end: The ending page number of the cited content within
            the document. Minimum value of 0.
    """

    documentIndex: int
    start: int
    end: int


class SearchResultLocation(TypedDict, total=False):
    """Specifies a search result location within the content array.

    Provides positioning information for cited content using search result
    index and block positions.

    Attributes:
        searchResultIndex: The index of the search result content block where
            the cited content is found. Minimum value of 0.
        start: The starting position in the content array where the cited
            content begins. Minimum value of 0.
        end: The ending position in the content array where the cited
            content ends. Minimum value of 0.
    """

    searchResultIndex: int
    start: int
    end: int


class WebLocation(TypedDict, total=False):
    """Provides the URL and domain information for a cited website.

    Contains information about the website that was cited when performing
    a web search.

    Attributes:
        url: The URL that was cited when performing a web search.
        domain: The domain that was cited when performing a web search.
    """

    url: str
    domain: str


# Tagged union type aliases following the ToolChoice pattern
DocumentCharLocationDict = dict[Literal["documentChar"], DocumentCharLocation]
DocumentPageLocationDict = dict[Literal["documentPage"], DocumentPageLocation]
DocumentChunkLocationDict = dict[Literal["documentChunk"], DocumentChunkLocation]
SearchResultLocationDict = dict[Literal["searchResultLocation"], SearchResultLocation]
WebLocationDict = dict[Literal["web"], WebLocation]

# Union type for citation locations - tagged union format matching AWS Bedrock API
CitationLocation = (
    DocumentCharLocationDict
    | DocumentPageLocationDict
    | DocumentChunkLocationDict
    | SearchResultLocationDict
    | WebLocationDict
)


class CitationSourceContent(TypedDict, total=False):
    """Contains the actual text content from a source document.

    Contains the actual text content from a source document that is being
    cited or referenced in the model's response.

    Note:
        This is a UNION type, so only one of the members can be specified.

    Attributes:
        text: The text content from the source document that is being cited.
    """

    text: str


class CitationGeneratedContent(TypedDict, total=False):
    """Contains the generated text content that corresponds to a citation.

    Contains the generated text content that corresponds to or is supported
    by a citation from a source document.

    Note:
        This is a UNION type, so only one of the members can be specified.

    Attributes:
        text: The text content that was generated by the model and is
            supported by the associated citation.
    """

    text: str


class Citation(TypedDict, total=False):
    """Contains information about a citation that references a source document.

    Citations provide traceability between the model's generated response
    and the source documents that informed that response.

    Attributes:
        location: The precise location within the source document where the
            cited content can be found, including character positions, page
            numbers, or chunk identifiers.
        sourceContent: The specific content from the source document that was
            referenced or cited in the generated response.
        title: The title or identifier of the source document being cited.
    """

    location: CitationLocation
    sourceContent: list[CitationSourceContent]
    title: str


class CitationsContentBlock(TypedDict, total=False):
    """A content block containing generated text and associated citations.

    This block type is returned when document citations are enabled, providing
    traceability between the generated content and the source documents that
    informed the response.

    Attributes:
        citations: An array of citations that reference the source documents
            used to generate the associated content.
        content: The generated content that is supported by the associated
            citations.
    """

    citations: list[Citation]
    content: list[CitationGeneratedContent]
