"""
  Embedding service (FastAPI).
  

  API (simple list-in, list-out; aligned by index):
  - POST /embed/text   body: ["text1", "text2", ...] -> [[...], ...]
  - POST /embed/image  body: ["url_or_path1", ...]  -> [[...], ...]

  """
  

  import logging

  import os

  import threading

  import time
  from collections import deque
  from dataclasses import dataclass

  from typing import Any, Dict, List, Optional
  
  import numpy as np

  from fastapi import FastAPI, HTTPException

  
  from embeddings.config import CONFIG

  from embeddings.protocols import ImageEncoderProtocol

  from config.services_config import get_embedding_backend_config

  logger = logging.getLogger(__name__)

  app = FastAPI(title="saas-search Embedding Service", version="1.0.0")

  # Models are loaded at startup, not lazily

  _text_model: Optional[Any] = None

  _image_model: Optional[ImageEncoderProtocol] = None

  _text_backend_name: str = ""

  open_text_model = os.getenv("EMBEDDING_ENABLE_TEXT_MODEL", "true").lower() in ("1", "true", "yes")
  open_image_model = os.getenv("EMBEDDING_ENABLE_IMAGE_MODEL", "true").lower() in ("1", "true", "yes")

  
  _text_encode_lock = threading.Lock()
  _image_encode_lock = threading.Lock()
  
  

  @dataclass
  class _SingleTextTask:
      text: str
      normalize: bool
      created_at: float
      done: threading.Event
      result: Optional[List[float]] = None
      error: Optional[Exception] = None
  
  
  _text_single_queue: "deque[_SingleTextTask]" = deque()
  _text_single_queue_cv = threading.Condition()
  _text_batch_worker: Optional[threading.Thread] = None
  _text_batch_worker_stop = False
  _TEXT_MICROBATCH_WINDOW_SEC = max(
      0.0, float(os.getenv("TEXT_MICROBATCH_WINDOW_MS", "4")) / 1000.0
  )
  _TEXT_REQUEST_TIMEOUT_SEC = max(
      1.0, float(os.getenv("TEXT_REQUEST_TIMEOUT_SEC", "30"))
  )
  
  
  def _encode_local_st(texts: List[str], normalize_embeddings: bool) -> Any:
      with _text_encode_lock:
          return _text_model.encode_batch(
              texts,
              batch_size=int(CONFIG.TEXT_BATCH_SIZE),
              device=CONFIG.TEXT_DEVICE,
              normalize_embeddings=normalize_embeddings,
          )
  
  
  def _start_text_batch_worker() -> None:
      global _text_batch_worker, _text_batch_worker_stop
      if _text_batch_worker is not None and _text_batch_worker.is_alive():
          return
      _text_batch_worker_stop = False
      _text_batch_worker = threading.Thread(
          target=_text_batch_worker_loop,
          name="embed-text-microbatch-worker",
          daemon=True,
      )
      _text_batch_worker.start()
      logger.info(
          "Started local_st text micro-batch worker | window_ms=%.1f max_batch=%d",
          _TEXT_MICROBATCH_WINDOW_SEC * 1000.0,
          int(CONFIG.TEXT_BATCH_SIZE),
      )
  
  
  def _stop_text_batch_worker() -> None:
      global _text_batch_worker_stop
      with _text_single_queue_cv:
          _text_batch_worker_stop = True
          _text_single_queue_cv.notify_all()
  
  
  def _text_batch_worker_loop() -> None:
      max_batch = max(1, int(CONFIG.TEXT_BATCH_SIZE))
      while True:
          with _text_single_queue_cv:
              while not _text_single_queue and not _text_batch_worker_stop:
                  _text_single_queue_cv.wait()
              if _text_batch_worker_stop:
                  return
  
              batch: List[_SingleTextTask] = [_text_single_queue.popleft()]
              deadline = time.perf_counter() + _TEXT_MICROBATCH_WINDOW_SEC
  
              while len(batch) < max_batch:
                  remaining = deadline - time.perf_counter()
                  if remaining <= 0:
                      break
                  if not _text_single_queue:
                      _text_single_queue_cv.wait(timeout=remaining)
                      continue
                  while _text_single_queue and len(batch) < max_batch:
                      batch.append(_text_single_queue.popleft())
  
          try:
              embs = _encode_local_st([task.text for task in batch], normalize_embeddings=False)
              if embs is None or len(embs) != len(batch):
                  raise RuntimeError(
                      f"Text model response length mismatch in micro-batch: "
                      f"expected {len(batch)}, got {0 if embs is None else len(embs)}"
                  )
              for task, emb in zip(batch, embs):
                  vec = _as_list(emb, normalize=task.normalize)
                  if vec is None:
                      raise RuntimeError("Text model returned empty embedding in micro-batch")
                  task.result = vec
          except Exception as exc:
              for task in batch:
                  task.error = exc
          finally:
              for task in batch:
                  task.done.set()
  
  
  def _encode_single_text_with_microbatch(text: str, normalize: bool) -> List[float]:
      task = _SingleTextTask(
          text=text,
          normalize=normalize,
          created_at=time.perf_counter(),
          done=threading.Event(),
      )
      with _text_single_queue_cv:
          _text_single_queue.append(task)
          _text_single_queue_cv.notify()
  
      if not task.done.wait(timeout=_TEXT_REQUEST_TIMEOUT_SEC):
          with _text_single_queue_cv:
              try:
                  _text_single_queue.remove(task)
              except ValueError:
                  pass
          raise RuntimeError(
              f"Timed out waiting for text micro-batch worker ({_TEXT_REQUEST_TIMEOUT_SEC:.1f}s)"
          )
      if task.error is not None:
          raise task.error
      if task.result is None:
          raise RuntimeError("Text micro-batch worker returned empty result")
      return task.result
  
  

  @app.on_event("startup")
  def load_models():
      """Load models at service startup to avoid first-request latency."""

      global _text_model, _image_model, _text_backend_name

      logger.info("Loading embedding models at startup...")

      # Load text model

      if open_text_model:
          try:

              backend_name, backend_cfg = get_embedding_backend_config()
              _text_backend_name = backend_name
              if backend_name == "tei":
                  from embeddings.tei_model import TEITextModel
  
                  base_url = (
                      os.getenv("TEI_BASE_URL")
                      or backend_cfg.get("base_url")
                      or CONFIG.TEI_BASE_URL
                  )
                  timeout_sec = int(
                      os.getenv("TEI_TIMEOUT_SEC")
                      or backend_cfg.get("timeout_sec")
                      or CONFIG.TEI_TIMEOUT_SEC
                  )
                  logger.info("Loading text backend: tei (base_url=%s)", base_url)
                  _text_model = TEITextModel(
                      base_url=str(base_url),
                      timeout_sec=timeout_sec,
                  )
              elif backend_name == "local_st":
                  from embeddings.qwen3_model import Qwen3TextModel

                  model_id = (
                      os.getenv("TEXT_MODEL_ID")
                      or backend_cfg.get("model_id")
                      or CONFIG.TEXT_MODEL_ID
                  )
                  logger.info("Loading text backend: local_st (model=%s)", model_id)
                  _text_model = Qwen3TextModel(model_id=str(model_id))

                  _start_text_batch_worker()

              else:
                  raise ValueError(
                      f"Unsupported embedding backend: {backend_name}. "
                      "Supported: tei, local_st"
                  )
              logger.info("Text backend loaded successfully: %s", _text_backend_name)

          except Exception as e:
              logger.error(f"Failed to load text model: {e}", exc_info=True)
              raise
      

      # Load image model: clip-as-service (recommended) or local CN-CLIP

      if open_image_model:
          try:

              if CONFIG.USE_CLIP_AS_SERVICE:

                  from embeddings.clip_as_service_encoder import ClipAsServiceImageEncoder
  

                  logger.info(f"Loading image encoder via clip-as-service: {CONFIG.CLIP_AS_SERVICE_SERVER}")
                  _image_model = ClipAsServiceImageEncoder(
                      server=CONFIG.CLIP_AS_SERVICE_SERVER,
                      batch_size=CONFIG.IMAGE_BATCH_SIZE,
                  )
                  logger.info("Image model (clip-as-service) loaded successfully")
              else:

                  from embeddings.clip_model import ClipImageModel
  

                  logger.info(f"Loading local image model: {CONFIG.IMAGE_MODEL_NAME} (device: {CONFIG.IMAGE_DEVICE})")
                  _image_model = ClipImageModel(
                      model_name=CONFIG.IMAGE_MODEL_NAME,
                      device=CONFIG.IMAGE_DEVICE,
                  )
                  logger.info("Image model (local CN-CLIP) loaded successfully")

          except Exception as e:

              logger.error("Failed to load image model: %s", e, exc_info=True)
              raise

  
      logger.info("All embedding models loaded successfully, service ready")

  @app.on_event("shutdown")
  def stop_workers() -> None:
      _stop_text_batch_worker()
  
  

  def _normalize_vector(vec: np.ndarray) -> np.ndarray:
      norm = float(np.linalg.norm(vec))
      if not np.isfinite(norm) or norm <= 0.0:
          raise RuntimeError("Embedding vector has invalid norm (must be > 0)")
      return vec / norm
  
  
  def _as_list(embedding: Optional[np.ndarray], normalize: bool = False) -> Optional[List[float]]:

      if embedding is None:
          return None
      if not isinstance(embedding, np.ndarray):
          embedding = np.array(embedding, dtype=np.float32)
      if embedding.ndim != 1:
          embedding = embedding.reshape(-1)

      embedding = embedding.astype(np.float32, copy=False)
      if normalize:
          embedding = _normalize_vector(embedding).astype(np.float32, copy=False)
      return embedding.tolist()

  
  
  @app.get("/health")
  def health() -> Dict[str, Any]:

      """Health check endpoint. Returns status and model loading state."""
      return {
          "status": "ok",
          "text_model_loaded": _text_model is not None,

          "text_backend": _text_backend_name,

          "image_model_loaded": _image_model is not None,
      }

  
  
  @app.post("/embed/text")

  def embed_text(texts: List[str], normalize: Optional[bool] = None) -> List[Optional[List[float]]]:

      if _text_model is None:
          raise RuntimeError("Text model not loaded")

      effective_normalize = bool(CONFIG.TEXT_NORMALIZE_EMBEDDINGS) if normalize is None else bool(normalize)

      normalized: List[str] = []

      for i, t in enumerate(texts):

          if not isinstance(t, str):

              raise HTTPException(status_code=400, detail=f"Invalid text at index {i}: must be string")
          s = t.strip()
          if not s:
              raise HTTPException(status_code=400, detail=f"Invalid text at index {i}: empty string")
          normalized.append(s)
  

      t0 = time.perf_counter()

      try:

          # local_st backend uses in-process torch model, keep serialized encode for safety;
          # TEI backend is an HTTP client and supports concurrent requests.
          if _text_backend_name == "local_st":
              if len(normalized) == 1 and _text_batch_worker is not None:
                  out = [_encode_single_text_with_microbatch(normalized[0], normalize=effective_normalize)]
                  elapsed_ms = (time.perf_counter() - t0) * 1000.0
                  logger.info(
                      "embed_text done | backend=%s mode=microbatch-single inputs=%d normalize=%s elapsed_ms=%.2f",
                      _text_backend_name,
                      len(normalized),
                      effective_normalize,
                      elapsed_ms,
                  )
                  return out
              embs = _encode_local_st(normalized, normalize_embeddings=False)
          else:

              embs = _text_model.encode_batch(
                  normalized,
                  batch_size=int(CONFIG.TEXT_BATCH_SIZE),
                  device=CONFIG.TEXT_DEVICE,

                  normalize_embeddings=effective_normalize,

              )
      except Exception as e:
          logger.error("Text embedding backend failure: %s", e, exc_info=True)
          raise HTTPException(
              status_code=502,
              detail=f"Text embedding backend failure: {e}",
          ) from e

      if embs is None or len(embs) != len(normalized):
          raise RuntimeError(
              f"Text model response length mismatch: expected {len(normalized)}, "
              f"got {0 if embs is None else len(embs)}"
          )
      out: List[Optional[List[float]]] = []
      for i, emb in enumerate(embs):

          vec = _as_list(emb, normalize=effective_normalize)

          if vec is None:
              raise RuntimeError(f"Text model returned empty embedding for index {i}")
          out.append(vec)

      elapsed_ms = (time.perf_counter() - t0) * 1000.0
      logger.info(
          "embed_text done | backend=%s inputs=%d normalize=%s elapsed_ms=%.2f",
          _text_backend_name,
          len(normalized),
          effective_normalize,
          elapsed_ms,
      )

      return out
  
  
  @app.post("/embed/image")

  def embed_image(images: List[str], normalize: Optional[bool] = None) -> List[Optional[List[float]]]:

      if _image_model is None:

          raise RuntimeError("Image model not loaded")

      effective_normalize = bool(CONFIG.IMAGE_NORMALIZE_EMBEDDINGS) if normalize is None else bool(normalize)

      urls: List[str] = []

      for i, url_or_path in enumerate(images):

          if not isinstance(url_or_path, str):

              raise HTTPException(status_code=400, detail=f"Invalid image at index {i}: must be string URL/path")
          s = url_or_path.strip()
          if not s:
              raise HTTPException(status_code=400, detail=f"Invalid image at index {i}: empty URL/path")
          urls.append(s)

      with _image_encode_lock:

          vectors = _image_model.encode_image_urls(
              urls,
              batch_size=CONFIG.IMAGE_BATCH_SIZE,
              normalize_embeddings=effective_normalize,
          )

      if vectors is None or len(vectors) != len(urls):
          raise RuntimeError(
              f"Image model response length mismatch: expected {len(urls)}, "
              f"got {0 if vectors is None else len(vectors)}"
          )
      out: List[Optional[List[float]]] = []
      for i, vec in enumerate(vectors):

          out_vec = _as_list(vec, normalize=effective_normalize)

          if out_vec is None:
              raise RuntimeError(f"Image model returned empty embedding for index {i}")
          out.append(out_vec)

      return out