"""
Base RAG Ingestion class.

Provides abstract methods for:
- OCR
- Chunking
- Embedding
- Vector Store operations

Providers can inherit and override methods as needed.
"""

from __future__ import annotations

import base64
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast

import litellm
from litellm._logging import verbose_logger
from litellm._uuid import uuid4
from litellm.constants import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE
from litellm.llms.custom_httpx.http_handler import (
    get_async_httpx_client,
    httpxSpecialProvider,
)
from litellm.litellm_core_utils.url_utils import async_safe_get
from litellm.rag.ingestion.file_parsers import extract_text_from_pdf
from litellm.rag.text_splitters import RecursiveCharacterTextSplitter
from litellm.types.rag import RAGIngestOptions, RAGIngestResponse

if TYPE_CHECKING:
    from litellm import Router


class BaseRAGIngestion(ABC):
    """
    Base class for RAG ingestion.

    Providers should inherit from this class and override methods as needed.
    For example, OpenAI handles embedding internally when attaching files to
    vector stores, so it overrides the embedding step to be a no-op.
    """

    supports_existing_file_id: bool = False

    def __init__(
        self,
        ingest_options: RAGIngestOptions,
        router: Optional["Router"] = None,
    ):
        self.ingest_options = ingest_options
        self.router = router
        self.ingest_id = f"ingest_{uuid4()}"

        # Extract configs from options
        self.ocr_config = ingest_options.get("ocr")
        self.chunking_strategy: Dict[str, Any] = cast(
            Dict[str, Any],
            ingest_options.get("chunking_strategy") or {"type": "auto"},
        )
        self.embedding_config = ingest_options.get("embedding")
        self.vector_store_config: Dict[str, Any] = cast(
            Dict[str, Any], ingest_options.get("vector_store") or {}
        )
        self.ingest_name = ingest_options.get("name")

        # Load credentials from litellm_credential_name if provided in vector_store config
        self._load_credentials_from_config()

    def _load_credentials_from_config(self) -> None:
        """
        Load credentials from litellm_credential_name if provided in vector_store config.

        This allows users to specify a credential name in the vector_store config
        which will be resolved from litellm.credential_list. When a stored
        credential is used, its values take precedence over caller-supplied
        equivalents so endpoint and identity fields stay consistent with the
        credential definition.
        """
        from litellm.litellm_core_utils.credential_accessor import CredentialAccessor

        credential_name = self.vector_store_config.get("litellm_credential_name")
        if credential_name and litellm.credential_list:
            credential_values = CredentialAccessor.get_credential_values(
                credential_name
            )
            if not credential_values:
                return
            for key, value in credential_values.items():
                self.vector_store_config[key] = value
            for key in (
                "api_base",
                "aws_sts_endpoint",
                "aws_web_identity_token",
            ):
                if key in self.vector_store_config and key not in credential_values:
                    del self.vector_store_config[key]

    @property
    def custom_llm_provider(self) -> str:
        """Get the vector store provider."""
        return self.vector_store_config.get("custom_llm_provider", "openai")

    async def upload(
        self,
        file_data: Optional[Tuple[str, bytes, str]] = None,
        file_url: Optional[str] = None,
        file_id: Optional[str] = None,
    ) -> Tuple[Optional[str], Optional[bytes], Optional[str], Optional[str]]:
        """
        Upload / prepare file for ingestion.

        Args:
            file_data: Tuple of (filename, content_bytes, content_type)
            file_url: URL to fetch file from
            file_id: Existing file ID to use

        Returns:
            Tuple of (filename, file_content, content_type, existing_file_id)
        """
        if file_data:
            filename, file_content, content_type = file_data
            return filename, file_content, content_type, None

        if file_url:
            http_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.RAG)
            response = await async_safe_get(http_client, file_url)
            response.raise_for_status()
            file_content = response.content
            filename = file_url.split("/")[-1] or "document"
            content_type = response.headers.get(
                "content-type", "application/octet-stream"
            )
            return filename, file_content, content_type, None

        if file_id:
            return None, None, None, file_id

        raise ValueError("Must provide file_data, file_url, or file_id")

    async def ocr(
        self,
        file_content: Optional[bytes],
        content_type: Optional[str],
    ) -> Optional[str]:
        """
        Perform OCR on file content to extract text.

        Args:
            file_content: Raw file bytes
            content_type: MIME type of the file

        Returns:
            Extracted text or None if OCR not configured/needed
        """
        if not self.ocr_config or not file_content:
            return None

        ocr_model = self.ocr_config.get("model", "mistral/mistral-ocr-latest")

        # Determine document type
        if content_type and "image" in content_type:
            doc_type, url_key = "image_url", "image_url"
        else:
            doc_type, url_key = "document_url", "document_url"

        # Encode as base64 data URL
        b64_content = base64.b64encode(file_content).decode("utf-8")
        data_url = f"data:{content_type};base64,{b64_content}"

        # Use router if available
        if self.router is not None:
            ocr_response = await self.router.aocr(
                model=ocr_model,
                document={"type": doc_type, url_key: data_url},
            )
        else:
            ocr_response = await litellm.aocr(
                model=ocr_model,
                document={"type": doc_type, url_key: data_url},
            )

        # Extract text from pages
        if hasattr(ocr_response, "pages") and ocr_response.pages:  # type: ignore
            return "\n\n".join(
                page.markdown for page in ocr_response.pages if hasattr(page, "markdown")  # type: ignore
            )

        return None

    def chunk(
        self,
        text: Optional[str],
        file_content: Optional[bytes],
        ocr_was_used: bool,
    ) -> List[str]:
        """
        Split text into chunks using RecursiveCharacterTextSplitter.

        Args:
            text: Text from OCR (if used)
            file_content: Raw file content bytes
            ocr_was_used: Whether OCR was performed

        Returns:
            List of text chunks
        """
        # Get text to chunk
        text_to_chunk: Optional[str] = None
        if text:
            text_to_chunk = text
        elif file_content and not ocr_was_used:
            # Try UTF-8 decode first
            try:
                text_to_chunk = file_content.decode("utf-8")
            except UnicodeDecodeError:
                # Check if it's a PDF and try to extract text
                if file_content.startswith(b"%PDF"):
                    verbose_logger.debug("PDF detected, attempting text extraction")
                    text_to_chunk = extract_text_from_pdf(file_content)
                    if not text_to_chunk:
                        verbose_logger.debug(
                            "PDF text extraction failed. Install 'pypdf' or 'PyPDF2' for PDF support, "
                            "or enable OCR with a vision model."
                        )
                        return []
                else:
                    verbose_logger.debug("Binary file detected, skipping text chunking")
                    return []

        if not text_to_chunk:
            return []

        # Extract RecursiveCharacterTextSplitter args
        splitter_args = self.chunking_strategy or {}
        chunk_size = splitter_args.get("chunk_size", DEFAULT_CHUNK_SIZE)
        chunk_overlap = splitter_args.get("chunk_overlap", DEFAULT_CHUNK_OVERLAP)
        separators = splitter_args.get("separators", None)

        # Build splitter kwargs
        splitter_kwargs: Dict[str, Any] = {
            "chunk_size": chunk_size,
            "chunk_overlap": chunk_overlap,
        }
        if separators:
            splitter_kwargs["separators"] = separators

        text_splitter = RecursiveCharacterTextSplitter(**splitter_kwargs)
        return text_splitter.split_text(text_to_chunk)

    async def embed(
        self,
        chunks: List[str],
    ) -> Optional[List[List[float]]]:
        """
        Generate embeddings for text chunks.

        Args:
            chunks: List of text chunks

        Returns:
            List of embeddings or None
        """
        if not self.embedding_config or not chunks:
            return None

        embedding_model = self.embedding_config.get("model", "text-embedding-3-small")

        if self.router is not None:
            response = await self.router.aembedding(model=embedding_model, input=chunks)
        else:
            response = await litellm.aembedding(model=embedding_model, input=chunks)

        return [item["embedding"] for item in response.data]

    @abstractmethod
    async def store(
        self,
        file_content: Optional[bytes],
        filename: Optional[str],
        content_type: Optional[str],
        chunks: List[str],
        embeddings: Optional[List[List[float]]],
        existing_file_id: str | None = None,
    ) -> Tuple[Optional[str], Optional[str]]:
        """
        Store content in vector store.

        This method must be implemented by provider-specific subclasses.

        Args:
            file_content: Raw file bytes
            filename: Name of the file
            content_type: MIME type
            chunks: Text chunks (if chunking was done locally)
            embeddings: Embeddings (if embedding was done locally)
            existing_file_id: Provider file ID supplied by the caller, if any

        Returns:
            Tuple of (vector_store_id, file_id)
        """
        pass

    async def ingest(
        self,
        file_data: Optional[Tuple[str, bytes, str]] = None,
        file_url: Optional[str] = None,
        file_id: Optional[str] = None,
    ) -> RAGIngestResponse:
        """
        Execute the full ingestion pipeline.

        Args:
            file_data: Tuple of (filename, content_bytes, content_type)
            file_url: URL to fetch file from
            file_id: Existing file ID to use

        Returns:
            RAGIngestResponse with status and IDs

        Raises:
            ValueError: If no input source is provided
        """
        # Step 1: Upload (raises ValueError if no input provided)
        filename, file_content, content_type, existing_file_id = await self.upload(
            file_data=file_data,
            file_url=file_url,
            file_id=file_id,
        )

        try:
            if existing_file_id and not self.supports_existing_file_id:
                raise ValueError(
                    f"{self.__class__.__name__} does not support ingesting an existing file_id. "
                    "Upload file data or provide file_url instead."
                )

            # Step 2: OCR (optional)
            extracted_text = await self.ocr(
                file_content=file_content,
                content_type=content_type,
            )

            # Step 3: Chunking
            chunks = self.chunk(
                text=extracted_text,
                file_content=file_content,
                ocr_was_used=self.ocr_config is not None,
            )

            # Step 4: Embedding (optional - some providers handle this internally)
            embeddings = await self.embed(chunks=chunks)

            # Step 5: Store in vector store
            vector_store_id, result_file_id = await self.store(
                file_content=file_content,
                filename=filename,
                content_type=content_type,
                chunks=chunks,
                embeddings=embeddings,
                existing_file_id=existing_file_id,
            )

            return RAGIngestResponse(
                id=self.ingest_id,
                status="completed",
                vector_store_id=vector_store_id or "",
                file_id=result_file_id or existing_file_id,
            )

        except Exception as e:
            verbose_logger.exception(f"RAG Pipeline failed: {e}")
            return RAGIngestResponse(
                id=self.ingest_id,
                status="failed",
                vector_store_id="",
                file_id=None,
                error=str(e),
            )
