# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
# SPDX-License-Identifier: AGPL-3.0
"""Jina AI Embedder Implementation"""

from typing import Any, Dict, List, Optional

import openai

from openviking.models.embedder.base import (
    DenseEmbedderBase,
    EmbedResult,
)
from openviking.utils.async_client_cache import LoopScopedAsyncClientCache
from openviking_cli.utils import get_logger

logger = get_logger(__name__)

# Default dimensions for Jina embedding models
JINA_MODEL_DIMENSIONS = {
    "jina-embeddings-v5-text-small": 1024,  # 677M params, max seq 32768
    "jina-embeddings-v5-text-nano": 768,  # 239M params, max seq 8192
    "jina-code-embeddings-1.5b": 1024,  # code model, max seq 8192
    "jina-code-embeddings-0.5b": 768,  # code model, max seq 8192
}

DEFAULT_JINA_QUERY_TASK = "retrieval.query"
DEFAULT_JINA_DOCUMENT_TASK = "retrieval.passage"
DEFAULT_JINA_CODE_QUERY_TASK = "nl2code.query"
DEFAULT_JINA_CODE_DOCUMENT_TASK = "nl2code.passage"
_UNSET = object()


def _get_default_task_params(model_name: str) -> tuple[str, str]:
    """Return the default Jina task names for the selected model."""
    if model_name.startswith("jina-code-embeddings-"):
        return DEFAULT_JINA_CODE_QUERY_TASK, DEFAULT_JINA_CODE_DOCUMENT_TASK
    return DEFAULT_JINA_QUERY_TASK, DEFAULT_JINA_DOCUMENT_TASK


class JinaDenseEmbedder(DenseEmbedderBase):
    """Jina AI Dense Embedder Implementation

    Uses Jina AI embedding API via OpenAI-compatible client.
    Supports task-specific embeddings (non-symmetric) and Matryoshka dimension reduction.

    Jina models are non-symmetric by default and require the 'task' parameter to distinguish
    between query and document embeddings. This is different from official OpenAI models,
    which are symmetric and do not support the input_type parameter.

    Example:
        >>> # Query embedding
        >>> query_embedder = JinaDenseEmbedder(
        ...     model_name="jina-embeddings-v5-text-small",
        ...     api_key="jina_xxx",
        ...     dimension=512,
        ...     context="query"
        ... )
        >>> query_vector = query_embedder.embed("search query")
        >>> print(len(query_vector.dense_vector))
        512

        >>> # Document embedding
        >>> doc_embedder = JinaDenseEmbedder(
        ...     model_name="jina-embeddings-v5-text-small",
        ...     api_key="jina_xxx",
        ...     dimension=512,
        ...     context="document"
        ... )
        >>> doc_vector = doc_embedder.embed("document content")
    """

    def __init__(
        self,
        model_name: str = "jina-embeddings-v5-text-small",
        api_key: Optional[str] = None,
        api_base: Optional[str] = None,
        dimension: Optional[int] = None,
        query_param: Any = _UNSET,
        document_param: Any = _UNSET,
        late_chunking: Optional[bool] = None,
        config: Optional[Dict[str, Any]] = None,
        task: Optional[str] = None,
    ):
        """Initialize Jina AI Dense Embedder

        Args:
            model_name: Jina model name, defaults to jina-embeddings-v5-text-small
            api_key: API key, required
            api_base: API base URL, defaults to https://api.jina.ai/v1
            dimension: Dimension for Matryoshka reduction, optional
            query_param: Task value for query-side embeddings. Defaults to 'retrieval.query'.
                        Override for models with different task naming conventions.
            document_param: Task value for document-side embeddings. Defaults to
                           'retrieval.passage'. Override for models with different task
                           naming conventions.
            late_chunking: Enable late chunking via extra_body, optional
            config: Additional configuration dict

        Raises:
            ValueError: If api_key is not provided
        """
        super().__init__(model_name, config)
        self.provider = "jina"

        self.api_key = api_key
        self.api_base = api_base or "https://api.jina.ai/v1"
        self.dimension = dimension
        default_query_param, default_document_param = _get_default_task_params(model_name)
        if query_param is _UNSET:
            query_param = default_query_param
        if document_param is _UNSET:
            document_param = default_document_param
        self.query_param = query_param
        self.document_param = document_param
        self.late_chunking = late_chunking

        if not self.api_key:
            raise ValueError("api_key is required")

        # Initialize OpenAI-compatible client with Jina base URL
        self.client = openai.OpenAI(
            api_key=self.api_key,
            base_url=self.api_base,
        )
        self._async_client_cache = LoopScopedAsyncClientCache()

        # Determine dimension
        max_dim = JINA_MODEL_DIMENSIONS.get(model_name, 1024)
        if dimension is not None and dimension > max_dim:
            raise ValueError(
                f"Requested dimension {dimension} exceeds maximum {max_dim} for model '{model_name}'. "
                f"Jina models support Matryoshka dimension reduction up to {max_dim}."
            )
        self._dimension = dimension if dimension is not None else max_dim

    def _build_extra_body(self, is_query: bool = False) -> Optional[Dict[str, Any]]:
        """Build extra_body dict for Jina-specific parameters"""
        extra_body = {}
        task = None
        if is_query and self.query_param is not None:
            task = self.query_param
        elif not is_query and self.document_param is not None:
            task = self.document_param

        if task is not None:
            extra_body["task"] = task
        if self.late_chunking is not None:
            extra_body["late_chunking"] = self.late_chunking
        return extra_body if extra_body else None

    def _build_kwargs(self, text_input: str | List[str], is_query: bool = False) -> Dict[str, Any]:
        kwargs: Dict[str, Any] = {"input": text_input, "model": self.model_name}
        if self.dimension:
            kwargs["dimensions"] = self.dimension

        extra_body = self._build_extra_body(is_query=is_query)
        if extra_body:
            kwargs["extra_body"] = extra_body
        return kwargs

    def _get_async_client(self):
        return self._async_client_cache.get(
            lambda: openai.AsyncOpenAI(
                api_key=self.api_key,
                base_url=self.api_base,
            )
        )

    def _raise_task_error(self, error: openai.APIError) -> None:
        """Raise an actionable error if a 422 indicates an invalid task type."""
        if getattr(error, "status_code", None) == 422 and "task" in str(error.body):
            raise RuntimeError(
                f"Jina API rejected task type for model '{self.model_name}'. "
                f"This usually means the model requires a different task prefix. "
                f"Set 'query_param' and 'document_param' in your embedding config "
                f"to a valid task type for this model. API details: {error.message}"
            ) from error

    def embed(self, text: str, is_query: bool = False) -> EmbedResult:
        """Perform dense embedding on text

        Args:
            text: Input text
            is_query: Flag to indicate if this is a query embedding

        Returns:
            EmbedResult: Result containing only dense_vector

        Raises:
            RuntimeError: When API call fails
        """

        def _call() -> EmbedResult:
            response = self.client.embeddings.create(**self._build_kwargs(text, is_query=is_query))
            vector = response.data[0].embedding

            return EmbedResult(dense_vector=vector)

        try:
            result = self._run_with_retry(
                _call,
                logger=logger,
                operation_name="Jina embedding",
            )
            # Estimate token usage
            estimated_tokens = self._estimate_tokens(text)
            self.update_token_usage(
                model_name=self.model_name,
                provider="jina",
                prompt_tokens=estimated_tokens,
                completion_tokens=0,
            )
            return result
        except openai.APIError as e:
            self._raise_task_error(e)
            raise RuntimeError(f"Jina API error: {e.message}") from e
        except Exception as e:
            raise RuntimeError(f"Embedding failed: {str(e)}") from e

    async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult:
        client = self._get_async_client()

        async def _call() -> EmbedResult:
            response = await client.embeddings.create(**self._build_kwargs(text, is_query=is_query))
            return EmbedResult(dense_vector=response.data[0].embedding)

        try:
            result = await self._run_with_async_retry(
                _call,
                logger=logger,
                operation_name="Jina async embedding",
            )
            estimated_tokens = self._estimate_tokens(text)
            self.update_token_usage(
                model_name=self.model_name,
                provider="jina",
                prompt_tokens=estimated_tokens,
                completion_tokens=0,
            )
            return result
        except openai.APIError as e:
            self._raise_task_error(e)
            raise RuntimeError(f"Jina API error: {e.message}") from e
        except Exception as e:
            raise RuntimeError(f"Embedding failed: {str(e)}") from e

    def embed_batch(self, texts: List[str], is_query: bool = False) -> List[EmbedResult]:
        """Batch embedding (Jina native support)

        Args:
            texts: List of texts
            is_query: Flag to indicate if these are query embeddings

        Returns:
            List[EmbedResult]: List of embedding results

        Raises:
            RuntimeError: When API call fails
        """
        if not texts:
            return []

        def _call() -> List[EmbedResult]:
            response = self.client.embeddings.create(**self._build_kwargs(texts, is_query=is_query))

            return [EmbedResult(dense_vector=item.embedding) for item in response.data]

        try:
            results = self._run_with_retry(
                _call,
                logger=logger,
                operation_name="Jina batch embedding",
            )
            # Estimate token usage for batch
            total_tokens = sum(self._estimate_tokens(text) for text in texts)
            self.update_token_usage(
                model_name=self.model_name,
                provider="jina",
                prompt_tokens=total_tokens,
                completion_tokens=0,
            )
            return results
        except openai.APIError as e:
            self._raise_task_error(e)
            raise RuntimeError(f"Jina API error: {e.message}") from e
        except Exception as e:
            raise RuntimeError(f"Batch embedding failed: {str(e)}") from e

    async def embed_batch_async(
        self, texts: List[str], is_query: bool = False
    ) -> List[EmbedResult]:
        if not texts:
            return []

        client = self._get_async_client()

        async def _call() -> List[EmbedResult]:
            response = await client.embeddings.create(
                **self._build_kwargs(texts, is_query=is_query)
            )
            return [EmbedResult(dense_vector=item.embedding) for item in response.data]

        try:
            results = await self._run_with_async_retry(
                _call,
                logger=logger,
                operation_name="Jina async batch embedding",
            )
            total_tokens = sum(self._estimate_tokens(text) for text in texts)
            self.update_token_usage(
                model_name=self.model_name,
                provider="jina",
                prompt_tokens=total_tokens,
                completion_tokens=0,
            )
            return results
        except openai.APIError as e:
            self._raise_task_error(e)
            raise RuntimeError(f"Jina API error: {e.message}") from e
        except Exception as e:
            raise RuntimeError(f"Batch embedding failed: {str(e)}") from e

    def get_dimension(self) -> int:
        """Get embedding dimension

        Returns:
            int: Vector dimension
        """
        return self._dimension
