"""TEI text embedding backend client."""
  
  from __future__ import annotations
  

  import logging

  from typing import Any, List, Union

  
  import numpy as np
  import requests
  

  logger = logging.getLogger(__name__)
  

  
  class TEITextModel:
      """
      Text embedding backend implemented via Hugging Face TEI HTTP API.
  
      Expected TEI endpoint:
        POST {base_url}/embed
        body: {"inputs": ["text1", "text2", ...]}
        response: [[...], [...], ...]
      """
  

      def __init__(self, base_url: str, timeout_sec: int = 60, max_client_batch_size: int = 24):

          if not base_url or not str(base_url).strip():
              raise ValueError("TEI base_url must not be empty")
          self.base_url = str(base_url).rstrip("/")
          self.endpoint = f"{self.base_url}/embed"
          self.timeout_sec = int(timeout_sec)

          self.max_client_batch_size = max(1, int(max_client_batch_size))

          self._health_check()
  
      def _health_check(self) -> None:
          health_url = f"{self.base_url}/health"
          response = requests.get(health_url, timeout=5)
          response.raise_for_status()

          # Probe one tiny embedding at startup so runtime requests do not fail later
          # with opaque "Invalid TEI embedding" errors.
          probe_resp = requests.post(
              self.endpoint,
              json={"inputs": ["health check"]},
              timeout=min(self.timeout_sec, 15),
          )
          probe_resp.raise_for_status()
          self._parse_payload(probe_resp.json(), expected_len=1)

  
      @staticmethod
      def _normalize(embedding: np.ndarray) -> np.ndarray:
          norm = np.linalg.norm(embedding)
          if norm <= 0:
              raise RuntimeError("TEI returned zero-norm embedding")
          return embedding / norm
  
      def encode(
          self,
          sentences: Union[str, List[str]],
          normalize_embeddings: bool = True,
          device: str = "cuda",
          batch_size: int = 32,
      ) -> np.ndarray:

          """
          Encode a single sentence or a list of sentences.

          TEI HTTP 后端天然是批量接口，这里统一通过 encode 处理单条和批量输入，
          不再额外暴露 encode_batch。
          """
  
          if isinstance(sentences, str):
              texts: List[str] = [sentences]
          else:
              texts = sentences

  
          if texts is None or len(texts) == 0:
              return np.array([], dtype=object)
          for i, t in enumerate(texts):
              if not isinstance(t, str) or not t.strip():
                  raise ValueError(f"Invalid input text at index {i}: {t!r}")
  

          if len(texts) > self.max_client_batch_size:
              logger.info(
                  "TEI batch split | total_inputs=%d chunk_size=%d chunks=%d",
                  len(texts),
                  self.max_client_batch_size,
                  (len(texts) + self.max_client_batch_size - 1) // self.max_client_batch_size,
              )
  
          vectors: List[np.ndarray] = []
          for start in range(0, len(texts), self.max_client_batch_size):
              batch = texts[start : start + self.max_client_batch_size]
              response = requests.post(
                  self.endpoint,
                  json={"inputs": batch},
                  timeout=self.timeout_sec,
              )
              response.raise_for_status()
              payload = response.json()
              parsed = self._parse_payload(payload, expected_len=len(batch))
              if normalize_embeddings:
                  parsed = [self._normalize(vec) for vec in parsed]
              vectors.extend(parsed)

          return np.array(vectors, dtype=object)

      def _parse_payload(self, payload: Any, expected_len: int) -> List[np.ndarray]:
          if not isinstance(payload, list) or len(payload) != expected_len:
              got = 0 if payload is None else (len(payload) if isinstance(payload, list) else "non-list")

              raise RuntimeError(

                  f"TEI response length mismatch: expected {expected_len}, got {got}. "
                  f"Response type={type(payload).__name__}"

              )
  
          vectors: List[np.ndarray] = []

          for i, item in enumerate(payload):
              emb = item.get("embedding") if isinstance(item, dict) else item
              try:
                  vec = np.asarray(emb, dtype=np.float32)
              except (TypeError, ValueError) as exc:
                  raise RuntimeError(
                      f"Invalid TEI embedding at index {i}: cannot convert to float array "
                      f"(item_type={type(item).__name__})"
                  ) from exc
  
              if vec.ndim != 1 or vec.size == 0:
                  raise RuntimeError(
                      f"Invalid TEI embedding at index {i}: shape={vec.shape}, size={vec.size}"
                  )
              if not np.isfinite(vec).all():
                  preview = vec[:8].tolist()
                  raise RuntimeError(
                      f"Invalid TEI embedding at index {i}: contains non-finite values, "
                      f"preview={preview}. This often indicates TEI backend/model runtime issues "
                      f"(for example an incompatible dtype or model config)."
                  )

              vectors.append(vec)

          return vectors