"""Embedding service"""

from __future__ import annotations

import json
import uuid
from abc import ABC, abstractmethod
from collections.abc import Callable
from typing import Any

import requests
from pydantic import BaseModel, Field

from ..core.config import get_embedding_config
from ..exceptions import RequestError
from ..utils.http_utils import build_base_url
from ..utils.log_utils import logger


class EmbeddingRequest(BaseModel):
    """Embedding 请求模型"""

    texts: list[str] = Field(..., description="待嵌入的文本列表")


class EmbeddingResult(BaseModel):
    """单条嵌入结果"""

    embedding: list[float] = Field(..., description="嵌入向量")
    index: int = Field(..., description="索引位置")


class EmbeddingResponse(BaseModel):
    """统一的 Embedding 响应模型"""

    results: list[EmbeddingResult] = Field(..., description="嵌入结果列表")


class EmbeddingConfig(BaseModel):
    """Embedding 服务配置"""

    model_name: str = Field(..., description="模型名称")
    dimension: int = Field(0, description="向量维度（0 表示使用模型默认维度）")
    api_key: str | None = Field(None, description="API密钥（如需要）")
    endpoints: list[str] = Field(..., description="API地址（用于负载均衡/容错）")
    extra: dict[str, Any] = Field(default_factory=dict, description="额外配置参数")


class EmbeddingService(ABC):
    """Embedding 服务基类"""

    def __init__(self, config: EmbeddingConfig) -> None:
        """初始化服务

        Args:
            config: 服务配置
        """
        self.config = config

    @abstractmethod
    def embed(self, texts: list[str]) -> EmbeddingResponse:
        """获取文本嵌入向量

        Args:
            texts: 待嵌入的文本列表

        Returns:
            嵌入响应
        """
        pass


# 服务注册表（支持通过名称创建服务）
_SERVICE_REGISTRY: dict[str, type[EmbeddingService]] = {}


def create_embedding_service(model_name: str, dimension: int = 0, **kwargs: Any) -> EmbeddingService:
    """创建 Embedding 服务

    Args:
        model_name: 模型名称（如 "bge-large-zh"）
        dimension: 向量维度（0 表示使用模型默认维度）
        **kwargs: 其他配置参数

    Returns:
        Embedding 服务实例

    Examples:
        >>> service = create_embedding_service("text-embed-7b", dimension=0, endpoints=["10.174.136.82:8691"])
        >>> response = service.embed(["你好", "世界"])
    """
    service_cls = _SERVICE_REGISTRY.get(model_name.strip())
    if not service_cls:
        raise ValueError(f"未知的模型: {model_name.strip()}. 支持的模型: {list(_SERVICE_REGISTRY.keys())}")

    need_endpoints = "endpoints" not in kwargs or not kwargs["endpoints"]
    should_backfill_api_key = "api_key" not in kwargs or kwargs["api_key"] is None

    cfg = None
    if need_endpoints or should_backfill_api_key:
        cfg = get_embedding_config()

    # if no endpoints args,get endpoints from config
    if need_endpoints:
        if cfg is None:
            raise RuntimeError("require embedding service config")
        # verify base name
        base_name = cfg.baseModelName
        if base_name != model_name:
            raise ValueError(f"模型名称与配置base_name不匹配: {model_name} != {base_name}")
        endpoint = cfg.endpoint
        kwargs["endpoints"] = [endpoint]

    if should_backfill_api_key:
        if cfg is not None and cfg.baseModelName == model_name:
            config_api_key = cfg.apiKey
            if config_api_key is not None:
                kwargs["api_key"] = config_api_key.strip() or None

    config = EmbeddingConfig(model_name=model_name.strip(), dimension=dimension, **kwargs)
    return service_cls(config)


def register_embedding_service(name: str) -> Callable[[type[EmbeddingService]], type[EmbeddingService]]:
    """注册自定义 Embedding 服务的装饰器

    Examples:
        >>> @register_embedding_service("my-model")
        >>> class MyEmbedding(EmbeddingService):
        ...     def embed(self, texts): ...
    """

    def decorator(cls: type[EmbeddingService]) -> type[EmbeddingService]:
        _SERVICE_REGISTRY[name.strip()] = cls
        return cls

    return decorator


@register_embedding_service("text-embed_7b_bf16")
class TextEmbed7B(EmbeddingService):
    """text-embed_7b_bf16 嵌入服务

    使用 Triton Inference Server 的 /v2/models/text_emb/infer 接口
    """

    def __init__(self, config: EmbeddingConfig) -> None:
        super().__init__(config)
        # 超时设置（秒）
        self.timeout_seconds = 30
        # 获取可用的端点列表
        self.endpoints = config.endpoints or []

        if not self.endpoints:
            raise ValueError("必须提供至少一个端点地址（endpoint 或 endpoints）")

    def _call_embedding_single(self, endpoint: str, text: str) -> list[float]:
        """调用单个端点获取单条文本的 embedding

        Args:
            endpoint: API端点地址
            text: 待嵌入的文本

        Returns:
            嵌入向量

        Raises:
            RuntimeError: 当API调用失败或响应格式错误时
        """
        url = f"{build_base_url(endpoint)}/v2/models/text_emb/infer"

        request_id = str(uuid.uuid4())
        logger.debug(f"Calling embedding endpoint: {url} with request_id: {request_id}")

        # 构造请求元数据
        meta_json = [
            {
                "request_id": request_id,
                "questions": [text],
                "algorithm": "text_emb",
            }
        ]
        json_str = json.dumps(meta_json, ensure_ascii=False)
        meta_bytes = json_str.encode("utf-8")

        # 构造请求
        inference_request = {
            "id": request_id,
            "inputs": [
                {
                    "name": "meta_json",
                    "shape": [1, len(meta_bytes)],
                    "datatype": "UINT8",
                    "data": [list(meta_bytes)],
                }
            ],
        }
        headers = {"Content-Type": "application/json"}

        # 发送请求
        response = requests.post(
            url,
            headers=headers,
            data=json.dumps(inference_request, ensure_ascii=False).encode("utf-8"),
            timeout=self.timeout_seconds,
        )
        response.encoding = "utf-8"

        # 检查 HTTP 状态码
        if response.status_code // 100 != 2:
            body = (response.text or "")[:200]
            raise RuntimeError(f"HTTP {response.status_code}: url={url} request_id={request_id} body={body!r}")

        # 解析 JSON 响应
        try:
            obj = response.json()
        except Exception as e:
            content_type = response.headers.get("content-type", "")
            body = (response.text or "")[:200]
            raise RuntimeError(
                f"Invalid JSON response: {e}; url={url}; request_id={request_id}; "
                f"status={response.status_code}; content-type={content_type}; body={body!r}"
            )

        # 解析 outputs
        outputs = obj.get("outputs")
        if not isinstance(outputs, list) or len(outputs) == 0:
            raise RuntimeError(f"missing outputs: url={url} request_id={request_id} resp={str(obj)[:200]}")

        output0 = outputs[0]
        if not isinstance(output0, dict):
            raise RuntimeError(f"invalid outputs[0]: url={url} request_id={request_id} resp={str(obj)[:200]}")

        output_data = output0.get("data")
        if not isinstance(output_data, list):
            raise RuntimeError(f"missing outputs[0].data: url={url} request_id={request_id} resp={str(obj)[:200]}")

        # 解码 UINT8 数据为 JSON
        try:
            json_payload = bytearray([v for v in output_data if isinstance(v, int) and 0 <= v <= 255]).decode(
                "utf-8", errors="strict"
            )
            ret_json = json.loads(json_payload)
        except Exception as e:
            raise RuntimeError(f"decode outputs[0].data failed: {e}; url={url}; request_id={request_id}")

        # 提取 embedding 向量
        if isinstance(ret_json, list):
            for item in ret_json:
                if not isinstance(item, dict):
                    continue

                # 尝试从 features 字段提取
                features = item.get("features")
                if isinstance(features, list) and len(features) > 0 and isinstance(features[0], list):
                    result: list[float] = features[0]
                    if result:
                        return result

                # 尝试从 embedding 字段提取
                emb = item.get("embedding")
                if isinstance(emb, list):
                    result = emb
                    if result:
                        return result

        raise RuntimeError(f"missing embedding: url={url} request_id={request_id} resp={str(ret_json)[:200]}")

    def embed(self, texts: list[str]) -> EmbeddingResponse:
        """获取文本嵌入向量（支持多端点容错）

        Args:
            texts: 待嵌入的文本列表

        Returns:
            嵌入响应

        Raises:
            RuntimeError: 当所有端点都失败时
        """
        results = []
        last_error: Exception | None = None

        for idx, text in enumerate(texts):
            # 尝试所有端点
            embedding: list[float] = []
            for endpoint in self.endpoints:
                try:
                    embedding = self._call_embedding_single(endpoint, text)
                    if embedding:  # 成功获取向量
                        break
                except Exception as e:
                    last_error = e
                    # 继续尝试下一个端点
                    continue

            results.append(EmbeddingResult(embedding=embedding, index=idx))

        # 如果所有文本都失败，抛出最后一个错误
        if all(len(r.embedding) == 0 for r in results) and last_error:
            raise RuntimeError(f"所有端点都失败: {last_error}")

        return EmbeddingResponse(results=results)


# bge-large-en 和 bge-large-zh 都基于千帆平台 API 接口
@register_embedding_service("bge-large-en")
@register_embedding_service("bge-large-zh")
class QianFanEmbedding(EmbeddingService):
    """千帆平台API接口嵌入服务"""

    def __init__(self, config: EmbeddingConfig) -> None:
        super().__init__(config)
        # 超时设置（秒）
        self.timeout_seconds = 30
        # 获取可用的端点列表
        self.endpoints = config.endpoints or []

    def _call_embedding_batch(self, endpoint: str, texts: list[str]) -> list[EmbeddingResult]:
        """调用单个端点获取多条文本的 embedding

        Args:
            endpoint: API端点地址
            texts: 待嵌入的文本列表

        Returns:
            嵌入向量

        Raises:
            RuntimeError: 当API调用失败或响应格式错误时
        """
        url = f"{build_base_url(endpoint)}/v2/embeddings"

        payload = json.dumps(
            {
                "model": self.config.model_name,
                "input": texts,
            }
        )
        headers = {"Content-Type": "application/json"}
        if self.config.api_key:
            headers["Authorization"] = f"Bearer {self.config.api_key}"

        response = requests.post(
            url,
            headers=headers,
            data=payload,
            timeout=self.timeout_seconds,
        )

        # 检查 HTTP 状态码
        if response.status_code // 100 != 2:
            raise RequestError(f"Request embedding model '{self.config.model_name}' failed")

        try:
            obj = response.json()
            # 解析并返回嵌入向量
            embeddings = obj.get("data", [])
            if isinstance(embeddings, list) and len(embeddings) > 0:
                results = []
                for idx, item in enumerate(embeddings):
                    vector = item.get("embedding")
                    if isinstance(vector, list):
                        results.append(EmbeddingResult(embedding=vector, index=idx))
                    else:
                        results.append(EmbeddingResult(embedding=[], index=idx))
                return results
            else:
                raise RequestError("Invalid embedding response format")
        except Exception:
            raise RequestError(f"Request embedding model '{self.config.model_name}' failed.")

    def embed(self, texts: list[str]) -> EmbeddingResponse:
        """获取文本嵌入向量（支持多端点容错）

        Args:
            texts: 待嵌入的文本列表

        Returns:
            嵌入响应

        Raises:
            RuntimeError: 当所有端点都失败时
        """
        last_error: Exception | None = None

        # 尝试所有端点
        embeddings: list[EmbeddingResult] = []
        for endpoint in self.endpoints:
            try:
                embeddings = self._call_embedding_batch(endpoint, texts)
                if embeddings:  # 成功获取向量
                    break
            except Exception as e:
                last_error = e
                # 继续尝试下一个端点
                continue

        # 如果所有文本都失败，抛出最后一个错误
        if all(len(r.embedding) == 0 for r in embeddings) and last_error:
            raise RequestError("all embedding endpoints failed.")

        return EmbeddingResponse(results=embeddings)
