"""
重排客户端：调用外部 BGE 重排服务，并对 ES 分数与重排分数进行融合。

流程：
1. 从 ES hits 构造用于重排的文档文本列表
2. POST 请求到重排服务 /rerank，获取每条文档的 relevance 分数
3. 将 ES 分数（归一化）与重排分数线性融合，写回 hit["_score"] 并重排序
"""

from typing import Dict, Any, List, Optional, Tuple
import logging

logger = logging.getLogger(__name__)

# 默认融合权重：ES 归一化分数权重、重排分数权重（相加为 1）
DEFAULT_WEIGHT_ES = 0.4
DEFAULT_WEIGHT_AI = 0.6
# 重排服务默认超时（文档较多时需更大，建议 config 中 timeout_sec 调大）
DEFAULT_TIMEOUT_SEC = 15.0


def build_docs_from_hits(
    es_hits: List[Dict[str, Any]],
    language: str = "zh",
) -> List[str]:
    """
    从 ES 命中结果构造重排服务所需的文档文本列表（与 hits 一一对应）。

    文本由 title、brief、description、vendor、category_path 等多语言字段拼接，
    按 language 优先选取对应语言；若无内容则用 spu_id 兜底。

    Args:
        es_hits: ES 返回的 hits 列表，每项含 _source
        language: 语言代码，如 "zh"、"en"

    Returns:
        与 es_hits 等长的字符串列表，用于 POST /rerank 的 docs
    """
    lang = (language or "zh").strip().lower()
    if lang not in ("zh", "en"):
        lang = "zh"

    def pick_lang_text(obj: Any) -> str:
        if obj is None:
            return ""
        if isinstance(obj, dict):
            return str(obj.get(lang) or obj.get("zh") or obj.get("en") or "").strip()
        return str(obj).strip()

    docs: List[str] = []
    for hit in es_hits:
        src = hit.get("_source") or {}
        parts: List[str] = []
        for key in ("title", "brief", "description", "vendor", "category_path"):
            parts.append(pick_lang_text(src.get(key)))
        text = " ".join(p for p in parts if p).strip()
        if not text:
            text = str(src.get("spu_id", ""))
        docs.append(text)
    return docs


def call_rerank_service(
    query: str,
    docs: List[str],
    service_url: str,
    timeout_sec: float = DEFAULT_TIMEOUT_SEC,
) -> Tuple[Optional[List[float]], Optional[Dict[str, Any]]]:
    """
    调用重排服务 POST /rerank，返回分数列表与 meta。

    Args:
        query: 搜索查询字符串
        docs: 文档文本列表（与 ES hits 顺序一致）
        service_url: 完整 URL，如 http://127.0.0.1:6007/rerank
        timeout_sec: 请求超时秒数

    Returns:
        (scores, meta)：成功时 scores 与 docs 等长，meta 为服务返回的 meta；
        失败时返回 (None, None)
    """
    if not docs:
        return [], {}
    try:
        import requests
        payload = {"query": (query or "").strip(), "docs": docs}
        response = requests.post(service_url, json=payload, timeout=timeout_sec)
        if response.status_code != 200:
            logger.warning(
                "Rerank service HTTP %s: %s",
                response.status_code,
                (response.text or "")[:200],
            )
            return None, None
        data = response.json()
        scores = data.get("scores")
        if not isinstance(scores, list):
            return None, None
        return scores, data.get("meta") or {}
    except (requests.exceptions.ReadTimeout, requests.exceptions.ConnectTimeout) as e:
        logger.warning(
            "Rerank request timed out after %.1fs (docs=%d); returning ES order. %s",
            timeout_sec, len(docs), e,
        )
        return None, None
    except Exception as e:
        logger.warning("Rerank request failed: %s", e, exc_info=True)
        return None, None


def fuse_scores_and_resort(
    es_hits: List[Dict[str, Any]],
    rerank_scores: List[float],
    weight_es: float = DEFAULT_WEIGHT_ES,
    weight_ai: float = DEFAULT_WEIGHT_AI,
) -> List[Dict[str, Any]]:
    """
    将 ES 分数与重排分数线性融合，写回每条 hit 的 _score，并按融合分数降序重排。

    对每条 hit 会写入：
    - _original_score: 原始 ES 分数
    - _ai_rerank_score: 重排服务返回的分数
    - _fused_score: 融合分数
    - _score: 置为融合分数（供后续 ResultFormatter 使用）

    Args:
        es_hits: ES hits 列表（会被原地修改）
        rerank_scores: 与 es_hits 等长的重排分数列表
        weight_es: ES 归一化分数权重
        weight_ai: 重排分数权重

    Returns:
        每条文档的融合调试信息列表，用于 debug_info
    """
    n = len(es_hits)
    if n == 0 or len(rerank_scores) != n:
        return []

    # 收集 ES 原始分数
    es_scores: List[float] = []
    for hit in es_hits:
        raw = hit.get("_score")
        try:
            es_scores.append(float(raw) if raw is not None else 0.0)
        except (TypeError, ValueError):
            es_scores.append(0.0)

    max_es = max(es_scores) if es_scores else 0.0
    fused_debug: List[Dict[str, Any]] = []

    for idx, hit in enumerate(es_hits):
        es_score = es_scores[idx]
        ai_score_raw = rerank_scores[idx]
        try:
            ai_score = float(ai_score_raw)
        except (TypeError, ValueError):
            ai_score = 0.0

        es_norm = (es_score / max_es) if max_es > 0 else 0.0
        fused = weight_es * es_norm + weight_ai * ai_score

        hit["_original_score"] = hit.get("_score")
        hit["_ai_rerank_score"] = ai_score
        hit["_fused_score"] = fused
        hit["_score"] = fused

        fused_debug.append({
            "doc_id": hit.get("_id"),
            "es_score": es_score,
            "es_score_norm": es_norm,
            "ai_rerank_score": ai_score,
            "fused_score": fused,
        })

    # 按融合分数降序重排
    es_hits.sort(
        key=lambda h: h.get("_fused_score", h.get("_score", 0.0)),
        reverse=True,
    )
    return fused_debug


def run_rerank(
    query: str,
    es_response: Dict[str, Any],
    language: str = "zh",
    service_url: Optional[str] = None,
    timeout_sec: float = DEFAULT_TIMEOUT_SEC,
    weight_es: float = DEFAULT_WEIGHT_ES,
    weight_ai: float = DEFAULT_WEIGHT_AI,
) -> Tuple[Dict[str, Any], Optional[Dict[str, Any]], List[Dict[str, Any]]]:
    """
    完整重排流程：从 es_response 取 hits -> 构造 docs -> 调服务 -> 融合分数并重排 -> 更新 max_score。

    Args:
        query: 搜索查询
        es_response: ES 原始响应（其中的 hits["hits"] 会被原地修改）
        language: 文档文本使用的语言
        service_url: 重排服务 URL，为 None 时使用默认 127.0.0.1:6007
        timeout_sec: 请求超时
        weight_es: ES 分数权重
        weight_ai: 重排分数权重

    Returns:
        (es_response, rerank_meta, fused_debug):
        - es_response: 已更新 hits 与 max_score 的响应（同一引用）
        - rerank_meta: 重排服务返回的 meta，失败时为 None
        - fused_debug: 每条文档的融合信息，供 debug 使用
    """
    try:
        from reranker.config import CONFIG as RERANKER_CONFIG
    except Exception:
        RERANKER_CONFIG = None

    url = service_url
    if not url and RERANKER_CONFIG is not None:
        url = f"http://127.0.0.1:{RERANKER_CONFIG.PORT}/rerank"
    if not url:
        url = "http://127.0.0.1:6007/rerank"

    hits = es_response.get("hits", {}).get("hits") or []
    if not hits:
        return es_response, None, []

    docs = build_docs_from_hits(hits, language=language)
    scores, meta = call_rerank_service(query, docs, url, timeout_sec=timeout_sec)

    if scores is None or len(scores) != len(hits):
        return es_response, None, []

    fused_debug = fuse_scores_and_resort(
        hits,
        scores,
        weight_es=weight_es,
        weight_ai=weight_ai,
    )

    # 更新 max_score 为融合后的最高分
    if hits:
        top = hits[0].get("_fused_score", hits[0].get("_score", 0.0)) or 0.0
        if "hits" in es_response:
            es_response["hits"]["max_score"] = top

    return es_response, meta, fused_debug