from __future__ import annotations

from typing import Any

from rag.ports import QueryEmbeddingUnavailable

BEDROCK_RUNTIME_SERVICE = "bedrock-runtime"
COHERE_EMBED_V4 = "cohere.embed-v4:0"

# Cohere embeds corpus records and queries into different subspaces. Sending
# the corpus value here raises no error — recall just drops silently — so this
# constant exists to make the asymmetry visible rather than incidental.
COHERE_QUERY_INPUT_TYPE = "search_query"


def _provider_error_types() -> tuple[type[BaseException], ...]:
    """botocore's error classes, or none when botocore is absent.

    Resolved lazily and tolerantly so a test injecting a stub client — and a
    deployment that never enables a cloud provider — neither imports botocore
    nor depends on it being installed.
    """
    try:
        from botocore.exceptions import BotoCoreError, ClientError
    except ImportError:
        return ()
    return (BotoCoreError, ClientError)


class BedrockCohereQueryEmbedder:
    """Embeds a query with the same model the collection was built from.

    A collection loaded with Cohere vectors and queried with any other embedder
    still returns results and still raises nothing — the hits are simply
    meaningless. That failure is silent, which is why the corpus manifest
    records `model_id` and why this adapter names the model explicitly.

    The request shape is duplicated from `ingestion/embed/bedrock_cohere.py`
    rather than imported: `ingestion` is a separate deployable and importing it
    here would couple the API to the batch pipeline. The duplication is one
    JSON body and is deliberate.
    """

    def __init__(
        self,
        dimensions: int,
        region: str = "us-east-1",
        client: Any | None = None,
        model_id: str = COHERE_EMBED_V4,
    ) -> None:
        if dimensions <= 0:
            raise ValueError("dimensions must be positive")
        self._dimensions = dimensions
        self._region = region
        self._client = client
        self._model_id = model_id

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

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

    def _runtime(self) -> Any:
        if self._client is None:
            import boto3
            from botocore.config import Config

            self._client = boto3.client(
                BEDROCK_RUNTIME_SERVICE,
                region_name=self._region,
                config=Config(
                    connect_timeout=10,
                    read_timeout=30,
                    # "adaptive" was tried 2026-08-07 and reverted same day — see
                    # the matching comment in adapters/bedrock_converse.py for the
                    # measured 1-5-minute regression it caused. "standard" kept.
                    retries={"max_attempts": 4, "mode": "standard"},
                ),
            )
        return self._client

    def embed_query(self, text: str) -> list[float]:
        import json

        try:
            response = self._runtime().invoke_model(
                modelId=self._model_id,
                body=json.dumps(
                    {
                        "texts": [text],
                        "input_type": COHERE_QUERY_INPUT_TYPE,
                        "embedding_types": ["float"],
                        "output_dimension": self._dimensions,
                    }
                ),
                accept="*/*",
                contentType="application/json",
            )
        except _provider_error_types() as error:
            raise QueryEmbeddingUnavailable(
                f"{self._model_id} could not be invoked: {type(error).__name__}"
            ) from error
        body = json.loads(response["body"].read())
        embeddings = body.get("embeddings")
        if isinstance(embeddings, dict):
            rows = embeddings.get("float")
        else:
            rows = embeddings
        if not rows:
            raise ValueError(
                f"{self._model_id} returned no float embedding; "
                f"response keys were {sorted(body)}"
            )
        values = list(rows[0])
        if len(values) != self._dimensions:
            raise ValueError(
                f"{self._model_id} returned {len(values)} dimensions, "
                f"expected {self._dimensions}"
            )
        return values