"""
Type definitions for RAG (Retrieval Augmented Generation) Ingest API.
"""

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

from pydantic import BaseModel, ConfigDict
from typing_extensions import TypedDict

from litellm.types.utils import ModelResponse


class RAGChunkingStrategy(TypedDict, total=False):
    """
    Chunking strategy config for RAG ingest using RecursiveCharacterTextSplitter.

    See: https://docs.langchain.com/oss/python/langchain/rag
    """

    chunk_size: int  # Maximum size of chunks (default: 1000)
    chunk_overlap: int  # Overlap between chunks (default: 200)
    separators: Optional[List[str]]  # Custom separators for splitting


class RAGIngestOCROptions(TypedDict, total=False):
    """OCR configuration for RAG ingest pipeline."""

    model: str  # e.g., "mistral/mistral-ocr-latest"


class RAGIngestEmbeddingOptions(TypedDict, total=False):
    """Embedding configuration for RAG ingest pipeline."""

    model: str  # e.g., "text-embedding-3-small"


class OpenAIVectorStoreOptions(TypedDict, total=False):
    """
    OpenAI vector store configuration.

    Example (auto-create):
        {"custom_llm_provider": "openai"}

    Example (use existing):
        {"custom_llm_provider": "openai", "vector_store_id": "vs_xxx"}

    Example (with credentials):
        {"custom_llm_provider": "openai", "litellm_credential_name": "my-openai-creds"}
    """

    custom_llm_provider: Literal["openai"]
    vector_store_id: Optional[str]  # Existing VS ID (auto-creates if not provided)
    ttl_days: Optional[int]  # Time-to-live in days for indexed content

    # Credentials (loaded from litellm.credential_list if litellm_credential_name is provided)
    litellm_credential_name: Optional[
        str
    ]  # Credential name to load from litellm.credential_list
    api_key: Optional[str]  # Direct API key (alternative to litellm_credential_name)
    api_base: Optional[str]  # Direct API base (alternative to litellm_credential_name)


class BedrockVectorStoreOptions(TypedDict, total=False):
    """
    Bedrock Knowledge Base configuration.

    Example (auto-create KB and all resources):
        {"custom_llm_provider": "bedrock"}

    Example (use existing KB):
        {"custom_llm_provider": "bedrock", "vector_store_id": "KB_ID"}

    Example (with credentials):
        {"custom_llm_provider": "bedrock", "litellm_credential_name": "my-aws-creds"}

    Auto-creation creates: S3 bucket, OpenSearch Serverless collection,
    IAM role, Knowledge Base, and Data Source.
    """

    custom_llm_provider: Literal["bedrock"]
    vector_store_id: Optional[str]  # Existing KB ID (auto-creates if not provided)

    # Bedrock-specific options
    s3_bucket: Optional[str]  # S3 bucket (auto-created if not provided)
    s3_prefix: Optional[str]  # S3 key prefix (default: "data/")
    embedding_model: Optional[
        str
    ]  # Embedding model (default: amazon.titan-embed-text-v2:0)
    data_source_id: Optional[str]  # For existing KB: override auto-detected DS
    wait_for_ingestion: Optional[
        bool
    ]  # Wait for completion (default: False - returns immediately)
    ingestion_timeout: Optional[
        int
    ]  # Timeout in seconds if wait_for_ingestion=True (default: 300)

    # Credentials (loaded from litellm.credential_list if litellm_credential_name is provided)
    litellm_credential_name: Optional[
        str
    ]  # Credential name to load from litellm.credential_list

    # AWS auth (uses BaseAWSLLM)
    aws_access_key_id: Optional[str]
    aws_secret_access_key: Optional[str]
    aws_session_token: Optional[str]
    aws_region_name: Optional[str]  # default: us-west-2
    aws_role_name: Optional[str]
    aws_session_name: Optional[str]
    aws_profile_name: Optional[str]
    aws_web_identity_token: Optional[str]
    aws_sts_endpoint: Optional[str]
    aws_external_id: Optional[str]


class VertexAIVectorStoreOptions(TypedDict, total=False):
    """
    Vertex AI RAG Engine configuration.

    Example (use existing corpus):
        {"custom_llm_provider": "vertex_ai", "vector_store_id": "CORPUS_ID", "gcs_bucket": "my-bucket"}

    Requires:
    - gcloud auth application-default login (for ADC authentication)
    - Files are uploaded to GCS via litellm.files.create_file, then imported into RAG corpus
    - GCS bucket must be provided via gcs_bucket or GCS_BUCKET_NAME env var
    """

    custom_llm_provider: Literal["vertex_ai"]
    vector_store_id: str  # RAG corpus ID (required for Vertex AI)

    # GCP config
    vertex_project: Optional[
        str
    ]  # GCP project ID (uses env VERTEXAI_PROJECT if not set)
    vertex_location: Optional[str]  # GCP region (default: us-central1)
    vertex_credentials: Optional[str]  # Path to credentials JSON (uses ADC if not set)
    gcs_bucket: Optional[
        str
    ]  # GCS bucket for file uploads (uses env GCS_BUCKET_NAME if not set)

    # Import settings
    wait_for_import: Optional[bool]  # Wait for import to complete (default: True)
    import_timeout: Optional[int]  # Timeout in seconds (default: 600)


class S3VectorsVectorStoreOptions(TypedDict, total=False):
    """
    AWS S3 Vectors configuration.

    Example (auto-create):
        {"custom_llm_provider": "s3_vectors", "vector_bucket_name": "my-embeddings"}

    Example (use existing):
        {"custom_llm_provider": "s3_vectors", "vector_bucket_name": "my-embeddings",
         "index_name": "my-index"}

    Example (with credentials):
        {"custom_llm_provider": "s3_vectors", "vector_bucket_name": "my-embeddings",
         "litellm_credential_name": "my-aws-creds"}

    Auto-creation creates: S3 vector bucket and vector index (if not provided).
    Embeddings are generated using LiteLLM's embedding API (supports any provider).
    """

    custom_llm_provider: Literal["s3_vectors"]
    vector_bucket_name: str  # Required - S3 vector bucket name
    index_name: Optional[str]  # Vector index name (auto-creates if not provided)

    # Index configuration (for auto-creation)
    dimension: Optional[
        int
    ]  # Vector dimension (auto-detected from embedding model, or default: 1024)
    distance_metric: Optional[Literal["cosine", "euclidean"]]  # Default: cosine
    non_filterable_metadata_keys: Optional[
        List[str]
    ]  # Keys excluded from filtering (e.g., ["source_text"])

    # Credentials (loaded from litellm.credential_list if litellm_credential_name is provided)
    litellm_credential_name: Optional[
        str
    ]  # Credential name to load from litellm.credential_list

    # AWS auth (uses BaseAWSLLM)
    aws_access_key_id: Optional[str]
    aws_secret_access_key: Optional[str]
    aws_session_token: Optional[str]
    aws_region_name: Optional[str]  # default: us-west-2
    aws_role_name: Optional[str]
    aws_session_name: Optional[str]
    aws_profile_name: Optional[str]
    aws_web_identity_token: Optional[str]
    aws_sts_endpoint: Optional[str]
    aws_external_id: Optional[str]


# Union type for vector store options
RAGIngestVectorStoreOptions = Union[
    OpenAIVectorStoreOptions,
    BedrockVectorStoreOptions,
    VertexAIVectorStoreOptions,
    S3VectorsVectorStoreOptions,
]


class RAGIngestOptions(TypedDict, total=False):
    """
    Combined options for RAG ingest pipeline.

    Unified interface - just specify custom_llm_provider:

    Example (OpenAI):
        from litellm.types.rag import RAGIngestOptions, OpenAIVectorStoreOptions

        options: RAGIngestOptions = {
            "vector_store": OpenAIVectorStoreOptions(
                custom_llm_provider="openai",
                vector_store_id="vs_xxx",  # optional
            )
        }

    Example (Bedrock):
        from litellm.types.rag import RAGIngestOptions, BedrockVectorStoreOptions

        options: RAGIngestOptions = {
            "vector_store": BedrockVectorStoreOptions(
                custom_llm_provider="bedrock",
                vector_store_id="KB_ID",  # optional - auto-creates if not provided
                wait_for_ingestion=True,
            )
        }
    """

    name: Optional[str]  # Optional pipeline name for logging
    ocr: Optional[RAGIngestOCROptions]  # Optional OCR step
    chunking_strategy: Optional[
        RAGChunkingStrategy
    ]  # RecursiveCharacterTextSplitter args
    embedding: Optional[RAGIngestEmbeddingOptions]  # Embedding model config
    vector_store: RAGIngestVectorStoreOptions  # OpenAI or Bedrock config


class RAGIngestResponse(TypedDict, total=False):
    """Response from RAG ingest API."""

    id: str  # Unique ingest job ID
    status: Literal["completed", "in_progress", "failed"]
    vector_store_id: str  # The vector store ID (created or existing)
    file_id: Optional[str]  # The file ID in the vector store
    error: Optional[str]  # Error message if status is "failed"


class RAGIngestRequest(BaseModel):
    """Request body for RAG ingest API (for validation)."""

    file_url: Optional[str] = None  # URL to fetch file from
    file_id: Optional[str] = None  # Existing file ID
    ingest_options: Dict[str, Any]  # RAGIngestOptions as dict for flexibility

    model_config = ConfigDict(extra="allow")  # Allow additional fields


class RAGRetrievalConfig(TypedDict, total=False):
    """Configuration for vector store retrieval."""

    vector_store_id: str
    custom_llm_provider: str
    top_k: int  # max results from vector store
    filters: Optional[Dict[str, Any]]  # optional - vector store filters


class RAGRerankConfig(TypedDict, total=False):
    """Configuration for reranking results."""

    enabled: bool
    model: str
    top_n: int  # final number of chunks after reranking
    return_documents: Optional[bool]


class RAGQueryRequest(BaseModel):
    """Request body for RAG query API."""

    model: str
    messages: List[Any]
    retrieval_config: RAGRetrievalConfig
    rerank: Optional[RAGRerankConfig] = None
    stream: Optional[bool] = False

    model_config = ConfigDict(extra="allow")


class RAGQueryResponse(ModelResponse):
    """Response from RAG query API."""

    pass
