"""Cohere Embed v4 (`cohere.embed-v4:0`).

Request/response shape taken from the AWS Bedrock user guide page "Cohere
Embed v4" (read 2026-08-03):

    request   {"input_type": "search_document|search_query|classification|
                              clustering",
               "texts": [str],            # max 96 per call
               "embedding_types": ["float"|"int8"|"uint8"|"binary"|"ubinary"],
               "output_dimension": 256|512|1024|1536,
               "truncate": "NONE|LEFT|RIGHT"}

The response has two documented shapes and this adapter accepts both. Asking
for one or more `embedding_types` returns
`{"response_type": "embeddings_by_type", "embeddings": {"float": [[...]]}}`;
omitting the field returns
`{"response_type": "embeddings_floats", "embeddings": [[...]]}`.

Three defaults are chosen rather than inherited:

- `output_dimension` is set explicitly. The documented default is 1536, and a
  collection built at one width cannot absorb vectors of another.
- `input_type` is derived from the caller's input kind. This is the model
  whose asymmetry the `ports` contract exists for: corpus records go in as
  `search_document`, queries as `search_query`.
- `truncate` is `NONE`, which makes an over-length input an error instead of a
  silently shortened one. A dosing section that lost its tail and embedded
  anyway is exactly the failure this project's rules are written against.
"""
from __future__ import annotations

from typing import Any, List, Sequence

from .bedrock_runtime import BedrockInvoker
from .ports import (
    INPUT_DOCUMENT,
    INPUT_QUERY,
    EmbeddingProvider,
    EmbeddingVector,
    text_digest,
)

MODEL_ID = "cohere.embed-v4:0"
PROVIDER_NAME = "cohere-v4"

SUPPORTED_DIMENSIONS = (256, 512, 1024, 1536)

# The documented per-request ceiling for `texts`.
MAX_TEXTS_PER_REQUEST = 96

COHERE_INPUT_TYPES = {
    INPUT_DOCUMENT: "search_document",
    INPUT_QUERY: "search_query",
}

# Cohere's docs do not state whether float vectors are unit-length, so this
# stays unset rather than being asserted either way.
NORMALIZED_UNKNOWN = None


class CohereEmbedV4(EmbeddingProvider):
    def __init__(
        self,
        invoker: BedrockInvoker,
        dimensions: int = 1024,
        truncate: str = "NONE",
        batch_size: int = MAX_TEXTS_PER_REQUEST,
    ):
        if dimensions not in SUPPORTED_DIMENSIONS:
            raise ValueError(
                f"{MODEL_ID} supports {SUPPORTED_DIMENSIONS}, got {dimensions}"
            )
        if not 1 <= batch_size <= MAX_TEXTS_PER_REQUEST:
            raise ValueError(
                f"batch_size must be 1..{MAX_TEXTS_PER_REQUEST}, got {batch_size}"
            )
        self._invoker = invoker
        self._dimensions = dimensions
        self._truncate = truncate
        self._batch_size = batch_size

    @property
    def name(self) -> str:
        return PROVIDER_NAME

    @property
    def model_id(self) -> str:
        return MODEL_ID

    @property
    def dimensions(self) -> int:
        return self._dimensions

    @property
    def max_batch_size(self) -> int:
        return self._batch_size

    def _embed_batch(
        self, texts: Sequence[str], input_kind: str
    ) -> List[EmbeddingVector]:
        body = self._invoker.invoke_json(
            MODEL_ID,
            {
                "texts": list(texts),
                "input_type": COHERE_INPUT_TYPES[input_kind],
                "embedding_types": ["float"],
                "output_dimension": self._dimensions,
                "truncate": self._truncate,
            },
            # The AWS code example for this model sends `*/*`.
            accept="*/*",
        )
        rows = _float_rows(body)
        if len(rows) != len(texts):
            raise ValueError(
                f"{MODEL_ID} returned {len(rows)} vectors for {len(texts)} texts"
            )

        vectors: List[EmbeddingVector] = []
        for text, values in zip(texts, rows, strict=True):
            self._check_dimensions(values)
            vectors.append(
                EmbeddingVector(
                    values=list(values),
                    text_sha256=text_digest(text),
                    provider=PROVIDER_NAME,
                    model_id=MODEL_ID,
                    dimensions=self._dimensions,
                    input_kind=input_kind,
                    normalized=NORMALIZED_UNKNOWN,
                )
            )
        return vectors


def _float_rows(body: Any) -> List[List[float]]:
    """Pull the float vectors out of either documented response shape."""
    embeddings = body.get("embeddings")
    if embeddings is None:
        raise ValueError(
            f"{MODEL_ID} response has no 'embeddings' field; "
            f"keys were {sorted(body)}"
        )
    if isinstance(embeddings, dict):
        rows = embeddings.get("float")
        if rows is None:
            raise ValueError(
                f"{MODEL_ID} returned no float embeddings; "
                f"types present: {sorted(embeddings)}"
            )
        return rows
    return embeddings