"""Rerank service implementation."""

from __future__ import annotations

import json
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_rerank_config
from ..exceptions import UnimplementedError
from ..utils.log_utils import logger


class RerankRequest(BaseModel):
    """Rerank 请求模型"""

    query: str = Field(..., description="查询文本")
    documents: list[str] = Field(..., description="待重排序的文档列表")
    top_k: int = Field(..., description="返回前K个结果")


class RerankResult(BaseModel):
    """单条重排序结果"""

    document: str = Field(..., description="文档内容")
    relevance_score: float = Field(..., description="相关性分数")
    index: int = Field(..., description="原始索引位置")


class RerankResponse(BaseModel):
    """统一的 Rerank 响应模型"""

    results: list[RerankResult] = Field(..., description="重排序结果列表")


class RerankConfig(BaseModel):
    """Rerank 服务配置"""

    model_name: str = Field(..., description="模型名称")
    api_key: str | None = Field(None, description="API密钥")
    endpoints: list[str] = Field(..., description="API地址（用于负载均衡/容错）")
    max_documents: int = Field(64, description="最大文档数量限制")
    max_top_k: int = Field(64, description="最大top_k限制")
    extra: dict[str, Any] = Field(default_factory=dict, description="额外配置参数")


class RerankService(ABC):
    """Rerank 服务基类"""

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

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

    @abstractmethod
    def rerank(self, query: str, documents: list[str], top_k: int) -> RerankResponse:
        """对文档进行重排序

        Args:
            query: 查询文本
            documents: 待重排序的文档列表
            top_k: 返回前K个结果

        Returns:
            重排序响应
        """
        raise UnimplementedError("子类必须实现 rerank 方法")

    def sort(self, original_obj: list[Any], rerank_result: list[RerankResult]) -> list[Any]:
        """根据rerank结果对原始对象进行排序

        Args:
            original_obj: 原始对象列表
            rerank_result: 重排序结果

        Returns:
            排序后的对象列表
        """
        sorted_obj = []
        for item in rerank_result:
            sorted_obj.append(original_obj[item.index])
        return sorted_obj


# 服务注册表
_SERVICE_REGISTRY: dict[str, type[RerankService]] = {}


def create_rerank_service(model_name: str, **kwargs: Any) -> RerankService:
    """创建 Rerank 服务

    Args:
        model_name: 模型名称（如 "bce-reranker-base"）
        **kwargs: 其他配置参数

    Returns:
        Rerank 服务实例

    Examples:
        >>> service = create_rerank_service("bce-reranker-base", api_key="bce-v3/ALTAK-XXX", endpoints=["https://qianfan.baidubce.com"])
        >>> response = service.rerank("查询", ["文档1", "文档2"], top_k=2)
    """
    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_rerank_config()

    # if no endpoints args,get endpoints from config
    if need_endpoints:
        if cfg is None:
            raise RuntimeError("require rerank 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 = RerankConfig(model_name=model_name.strip(), **kwargs)
    return service_cls(config)


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

    Examples:
        >>> @register_rerank_service("my-reranker")
        >>> class MyReranker(RerankService):
        ...     def rerank(self, query, documents, top_k): ...
    """

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

    return decorator


class QianFanReranker(RerankService):
    """千帆平台 BCE Reranker 服务"""

    def __init__(self, config: RerankConfig) -> None:
        super().__init__(config)
        self.timeout_seconds = 30
        self.endpoints = config.endpoints or []
        self.path = "/v2/rerank"

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

    def rerank(self, query: str, documents: list[str], top_k: int) -> RerankResponse:
        """对文档进行重排序

        Args:
            query: 查询文本
            documents: 待重排序的文档列表
            top_k: 返回前K个结果

        Returns:
            重排序响应

        Raises:
            ValueError: 当参数不合法时
            RuntimeError: 当API调用失败时
        """
        # 参数验证和调整
        top_k = self._validate_and_adjust_params(documents, top_k)

        # 尝试所有端点
        for endpoint in self.endpoints:
            try:
                return self._call_rerank_single(endpoint, query, documents, top_k)
            except Exception as e:
                logger.warning(f"Rerank endpoint {endpoint} failed: {e}")
                continue

        # 所有端点都失败
        raise RuntimeError("请求 Rerank 服务失败")

    def _validate_and_adjust_params(self, documents: list[str], top_k: int) -> int:
        """验证并调整参数

        Args:
            documents: 文档列表
            top_k: 期望返回的数量

        Returns:
            调整后的 top_k
        """
        if not documents:
            raise ValueError("documents 不能为空")

        # 调整 top_k
        if top_k > len(documents):
            logger.warning(f"top_k ({top_k}) 大于 documents 数量 ({len(documents)})，使用 documents 数量")
            top_k = len(documents)

        if top_k <= 0:
            if len(documents) > 0:
                top_k = len(documents)
                logger.warning(f"top_k 未设置或小于等于0，使用 documents 数量: {top_k}")
            else:
                raise ValueError("top_k 必须大于0")

        if top_k > self.config.max_top_k:
            logger.warning(f"top_k ({top_k}) 超过限制 ({self.config.max_top_k})，截断为 {self.config.max_top_k}")
            top_k = self.config.max_top_k

        if len(documents) > self.config.max_documents:
            logger.warning(f"documents 数量 ({len(documents)}) 超过限制 ({self.config.max_documents})，将截断")
            documents = documents[: self.config.max_documents]
        return top_k

    def _call_rerank_single(self, endpoint: str, query: str, documents: list[str], top_k: int) -> RerankResponse:
        """调用单个端点进行重排序

        Args:
            endpoint: API端点地址
            query: 查询文本
            documents: 待重排序的文档列表
            top_k: 返回前K个结果

        Returns:
            重排序响应

        Raises:
            RuntimeError: 当API调用失败时
        """
        # 截断文档列表
        documents = documents[: self.config.max_documents]

        url = f"{endpoint.rstrip('/')}{self.path}"
        logger.debug(f"Calling rerank endpoint: {url}")

        payload = json.dumps(
            {
                "model": self.config.model_name,
                "query": query,
                "documents": documents,
                "top_n": min(top_k, self.config.max_top_k),
            },
            ensure_ascii=False,
        )

        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.encode("utf-8"),
            timeout=self.timeout_seconds,
        )

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

        # 解析响应
        try:
            result = response.json()
        except Exception as e:
            body = (response.text or "")[:200]
            raise RuntimeError(f"Invalid JSON response: {e}; url={url}; body={body!r}")

        # 转换为标准响应格式
        results = result.get("results", [])
        rerank_results = [
            RerankResult(
                document=item.get("document", ""),
                relevance_score=item.get("relevance_score", 0.0),
                index=item.get("index", 0),
            )
            for item in results
        ]

        return RerankResponse(results=rerank_results)


@register_rerank_service("bce-reranker-base")
class LogicReranker(QianFanReranker):
    def __init__(self, config: RerankConfig) -> None:
        super().__init__(config)
        # Logic Reranker 使用不同的路径
        self.path = "/rerankers"


if __name__ == "__main__":
    # 使用新的接口
    service = create_rerank_service(
        "bce-reranker-base",
        api_key="bce-v3/ALTAK-xG0aAOmjm1C7tRrCK4CQd/4e1041ab70be617cd41058c3c02a5c890cbe2916",
        endpoints=["https://qianfan.baidubce.com"],
    )

    response = service.rerank(
        query="北京人工智能怎么样？",
        documents=[
            "致力于创建能够执行通常需要人类智能的任务的系统。",
            "机器学习、自然语言处理、计算机视觉等领域。",
            "北京天气。",
        ],
        top_k=2,
    )

    for result in response.results:
        print(f"Document: {result.document}")
        print(f"Score: {result.relevance_score}")
        print(f"Index: {result.index}")
        print("---")
