from typing import Any, Dict, List, Optional, Union

from typing_extensions import TypedDict


class RAGImage(TypedDict, total=False):
    full_folder_id: str
    file_path: str
    index: int
    file_data: str


# New in DSS 14
class SourceRecords(TypedDict, total=False):
    columns: Optional[List[str]]
    data: Optional[List[List[Dict]]]


# New in DSS 14
class ImageRef(TypedDict, total=False):
    folderId: Optional[str]
    path: Optional[str]


# New in DSS 14
class DocumentPageRange(TypedDict, total=False):
    start: Optional[int]
    end: Optional[int]


# New in DSS 14
class FileRef(TypedDict, total=False):
    folderId: Optional[str]
    path: Optional[str]
    pageRange: Optional[DocumentPageRange]


class Source(TypedDict, total=False):
    # We have similarities with DSS sources but we also have some
    # specific things related to the way we build Sources

    # Standard types
    # - FILE_BASED_DOCUMENT: means there will be a fileRef with the path of the reference document
    # - SIMPLE_DOCUMENT:     will still be used if the informations comes from a dataset embed recipe and
    #                        there's no specific document as source
    # - RECORDS
    # - GENERATED_SQL_QUERY: SQL query that will be rendered with syntax highlighting
    # - GENERATED_CYPHER_QUERY: Cypher query that will be rendered with syntax highlighting
    type: Optional[str]  # new type FILE_BASED_DOCUMENT

    # Generic usage
    metadata: Optional[Dict[str, Union[str, int, float, List[Any]]]]

    # For RECORDS
    records: Optional[SourceRecords]

    # For FILE_BASED_DOCUMENT
    fileRef: Optional[FileRef]  # new in DSS 14
    imageRefs: Optional[List[ImageRef]]  # new in DSS 14

    # For SIMPLE_DOCUMENT
    title: Optional[str]
    url: Optional[str]
    thumbnailInlineB64: Optional[str]  # new in DSS 14
    thumbnailURL: Optional[str]  # new in DSS 14
    thumbnailW: Optional[int]  # new in DSS 14
    thumbnailH: Optional[int]  # new in DSS 14

    # For SIMPLE_DOCUMENT and FILE_BASED_DOCUMENT
    textSnippet: Optional[str]
    markdownSnippet: Optional[str]  # new - to handle in front
    htmlSnippet: Optional[str]  # new - to handle in front
    jsonSnippet: Optional[str]  # new - to handle in front

    # For GENERATED_SQL_QUERY and GENERATED_CYPHER_QUERY (DSS 14 and earlier)
    performedQuery: Optional[str]

    # For GENERATED_SQL_QUERY and GENERATED_CYPHER_QUERY (DSS 15+)
    performedQueries: Optional[List[str]]

    images: Optional[List[RAGImage]]


# if you update it think to common.backend.schemas.common_schemas.py::SourceSchema
# but be careful with what should be reflected in the API and what not
class AggregatedToolSources(TypedDict, total=False):
    toolCallDescription: Optional[str]
    items: Optional[List[Source]]
    # Raw input only: identifies the originating tool call, found in hierarchy[].toolCallId
    # where hierarchy[].type == "TOOL". Not part of the upstream SourceSchema.
    hierarchy: Optional[List[Dict[str, Any]]]
    # Output only: threaded from hierarchy so the frontend can join a source with its
    # tool call (see get_vector_search_tool_calls) by id instead of by array position.
    toolCallId: Optional[str]
    # Raw input only: KB search tool calls performed by an agent called as a tool
    # (see AHDSSAgentTool.run in agent_connect_v2.py). tools_using_v2.py's generic
    # tool-call  only forwards the `sources` field across the tool-call
    # boundary, so this rides along here instead of a dedicated channel. Consumed and
    # stripped by sources_utils.extract_nested_vector_search_tool_calls; never reaches
    # the frontend as-is.
    nestedVectorSearchToolCalls: Optional[List[Dict[str, Any]]]
