"""
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
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 = True
open_image_model = True  # Enable image embedding when using clip-as-service

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


@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))
            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")


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)

    try:
        with _text_encode_lock:
            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)
    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