# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Embedding API protocol models for OpenAI and Cohere formats.

OpenAI: https://platform.openai.com/docs/api-reference/embeddings
Cohere: https://docs.cohere.com/reference/embed
"""

import builtins
import struct
import time
from collections.abc import Sequence
from typing import Annotated, Any, Literal, TypeAlias

import numpy as np
import pybase64 as base64
from pydantic import BaseModel, BeforeValidator, Field, model_validator

from vllm import PoolingParams
from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo
from vllm.utils import random_uuid

from ..base.protocol import (
    ChatRequestMixin,
    ChatRequestOptionsMixin,
    CompletionRequestMixin,
    EmbeddingTokenizeParamsMixin,
    EmbedRequestMixin,
    PoolingBasicRequestMixin,
    reject_removed_pooling_parameters,
)


class EmbeddingCompletionRequest(
    PoolingBasicRequestMixin,
    CompletionRequestMixin,
    EmbedRequestMixin,
    EmbeddingTokenizeParamsMixin,
):
    def to_pooling_params(self):
        return PoolingParams(
            task="embed",
            dimensions=self.dimensions,
            use_activation=self.use_activation,
        )


def _is_chat_message(value: Any) -> bool:
    return isinstance(value, dict) and isinstance(value.get("role"), str)


def _is_chat_messages(value: Any) -> bool:
    return (
        isinstance(value, list)
        and bool(value)
        and all(_is_chat_message(item) for item in value)
    )


def _is_batched_chat_messages(value: Any) -> bool:
    return (
        isinstance(value, list)
        and bool(value)
        and all(_is_chat_messages(item) for item in value)
    )


class EmbeddingChatRequest(
    PoolingBasicRequestMixin,
    ChatRequestMixin,
    EmbedRequestMixin,
    EmbeddingTokenizeParamsMixin,
):
    """OpenAI embeddings request with one top-level chat conversation."""

    def to_pooling_params(self):
        return PoolingParams(
            task="embed",
            dimensions=self.dimensions,
            use_activation=self.use_activation,
        )


class EmbeddingBatchChatRequest(
    PoolingBasicRequestMixin,
    ChatRequestOptionsMixin,
    EmbedRequestMixin,
    EmbeddingTokenizeParamsMixin,
):
    """OpenAI embeddings request with batched top-level chat conversations.

    Mirrors ``BatchChatCompletionRequest`` by keeping batched conversations in
    ``messages`` instead of introducing a separate batch-specific field.
    """

    messages: Sequence[
        Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]
    ] = Field(..., min_length=1)

    def to_pooling_params(self):
        return PoolingParams(
            task="embed",
            dimensions=self.dimensions,
            use_activation=self.use_activation,
        )


class EmbeddingChatInputRequest(
    EmbeddingChatRequest,
):
    """OpenAI embeddings request with one chat conversation in ``input``."""

    input: list[ChatCompletionMessageParam]

    @model_validator(mode="before")
    @classmethod
    def normalize_input_messages(cls, data):
        if not isinstance(data, dict):
            return data

        if "messages" in data or "input" not in data:
            return data

        input_data = data["input"]
        if not _is_chat_messages(input_data):
            return data

        normalized = dict(data)
        normalized["messages"] = input_data
        return normalized


class EmbeddingBatchChatInputRequest(EmbeddingBatchChatRequest):
    """OpenAI embeddings request with batched chat conversations in ``input``."""

    input: Sequence[
        Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]
    ] = Field(..., min_length=1)

    @model_validator(mode="before")
    @classmethod
    def normalize_input_messages(cls, data):
        if not isinstance(data, dict):
            return data

        if "messages" in data or "input" not in data:
            return data

        input_data = data["input"]
        if not _is_batched_chat_messages(input_data):
            return data

        normalized = dict(data)
        normalized["messages"] = input_data
        return normalized


EmbeddingRequest: TypeAlias = Annotated[
    EmbeddingCompletionRequest
    | EmbeddingChatRequest
    | EmbeddingBatchChatRequest
    | EmbeddingChatInputRequest
    | EmbeddingBatchChatInputRequest,
    BeforeValidator(reject_removed_pooling_parameters),
]


# ---------------------------------------------------------------------------
# OpenAI /v1/embeddings — response models
# ---------------------------------------------------------------------------


class EmbeddingResponseData(OpenAIBaseModel):
    index: int
    object: str = "embedding"
    embedding: list[float] | str


class EmbeddingResponse(OpenAIBaseModel):
    id: str = Field(default_factory=lambda: f"embd-{random_uuid()}")
    object: str = "list"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str | None = None
    data: list[EmbeddingResponseData]
    usage: UsageInfo


class EmbeddingBytesResponse(OpenAIBaseModel):
    content: list[bytes]
    headers: dict[str, str] | None = None
    media_type: str = "application/octet-stream"


# ---------------------------------------------------------------------------
# Cohere /v2/embed — request models
# ---------------------------------------------------------------------------

CohereEmbeddingType = Literal[
    "float",
    "binary",
    "ubinary",
    "base64",
]
CohereTruncate = Literal["NONE", "START", "END"]


class CohereEmbedContent(BaseModel):
    type: Literal["text", "image_url"]
    text: str | None = None
    image_url: dict[str, str] | None = None

    @model_validator(mode="after")
    def validate_content_payload(self):
        if self.type == "text":
            if self.text is None:
                raise ValueError("CohereEmbedContent with type='text' requires text")
        elif not self.image_url or not self.image_url.get("url"):
            raise ValueError(
                "CohereEmbedContent with type='image_url' requires image_url.url"
            )
        return self


class CohereEmbedInput(BaseModel):
    content: list[CohereEmbedContent]


class CohereEmbedRequest(BaseModel):
    model: str | None = None
    input_type: str | None = None
    texts: list[str] | None = None
    images: list[str] | None = None
    inputs: list[CohereEmbedInput] | None = None
    output_dimension: int | None = None
    embedding_types: list[CohereEmbeddingType] | None = None
    truncate: CohereTruncate = "END"
    max_tokens: int | None = None
    priority: int = 0

    @model_validator(mode="after")
    def validate_input_fields(self):
        input_fields = (self.texts, self.images, self.inputs)
        provided_fields = [field for field in input_fields if field is not None]
        if len(provided_fields) != 1 or not provided_fields[0]:
            raise ValueError(
                "Exactly one of texts, images, or inputs must be provided, "
                "and it must be non-empty"
            )
        return self


# ---------------------------------------------------------------------------
# Cohere /v2/embed — response models
# ---------------------------------------------------------------------------


class CohereApiVersion(BaseModel):
    version: str = "2"


class CohereBilledUnits(BaseModel):
    input_tokens: int | None = None
    image_tokens: int | None = None


class CohereMeta(BaseModel):
    api_version: CohereApiVersion = Field(default_factory=CohereApiVersion)
    billed_units: CohereBilledUnits | None = None


class CohereEmbedByTypeEmbeddings(BaseModel):
    # The field name ``float`` shadows the builtin type, so the annotation
    # must use ``builtins.float`` to avoid a self-referential type error.
    float: list[list[builtins.float]] | None = None
    binary: list[list[int]] | None = None
    ubinary: list[list[int]] | None = None
    base64: list[str] | None = None


class CohereEmbedResponse(BaseModel):
    id: str = Field(default_factory=lambda: f"embd-{random_uuid()}")
    embeddings: CohereEmbedByTypeEmbeddings
    texts: list[str] | None = None
    meta: CohereMeta | None = None
    response_type: Literal["embeddings_by_type"] = "embeddings_by_type"


# ---------------------------------------------------------------------------
# Cohere embedding type conversion helpers
# ---------------------------------------------------------------------------

_UNSIGNED_TO_SIGNED_DIFF = 1 << 7  # 128


def _pack_binary_embeddings(
    float_embeddings: list[list[float]],
    signed: bool,
) -> list[list[int]]:
    """Bit-pack float embeddings: positive -> 1, negative -> 0.

    Bits are packed MSB-first, eight per byte.
    """
    if not float_embeddings:
        return []

    array = np.asarray(float_embeddings, dtype=np.float64)
    if array.ndim != 2:
        raise ValueError(
            f"Expected a 2D batch of embeddings, but got {array.ndim}D input."
        )

    dim = array.shape[1]
    if dim % 8 != 0:
        raise ValueError(
            "Embedding dimension must be a multiple of 8 for binary "
            f"embedding types, but got {dim}."
        )

    packed = np.packbits(array >= 0, axis=-1)
    if signed:
        packed = packed.astype(np.int16) - _UNSIGNED_TO_SIGNED_DIFF
    return packed.tolist()


def _encode_base64_embeddings(
    float_embeddings: list[list[float]],
) -> list[str]:
    """Encode float embeddings as base64 (little-endian float32)."""
    result: list[str] = []
    for embedding in float_embeddings:
        buf = struct.pack(f"<{len(embedding)}f", *embedding)
        result.append(base64.b64encode(buf).decode("utf-8"))
    return result


def build_typed_embeddings(
    float_embeddings: list[list[float]],
    embedding_types: Sequence[str],
) -> CohereEmbedByTypeEmbeddings:
    """Convert float embeddings to all requested Cohere embedding types."""
    result = CohereEmbedByTypeEmbeddings()

    for emb_type in embedding_types:
        if emb_type == "float":
            result.float = float_embeddings
        elif emb_type == "binary":
            result.binary = _pack_binary_embeddings(float_embeddings, signed=True)
        elif emb_type == "ubinary":
            result.ubinary = _pack_binary_embeddings(float_embeddings, signed=False)
        elif emb_type == "base64":
            result.base64 = _encode_base64_embeddings(float_embeddings)

    return result