# requires: elasticsearch
"""Elasticsearch：中文分词、BM25 关键词检索、kNN 向量检索、RRF 混合检索、别名零停机重建索引。

运行：ES_URL=http://127.0.0.1:9200 python es_hybrid.py
本地测试用的单节点 ES 关闭了安全认证；生产环境必须开启认证和 TLS。没有设置环境变量时跳过。
向量用一个“字符二元组哈希”的玩具嵌入代替真实模型，只为演示流程；索引名带随机后缀，结束后删除。
"""
import hashlib
import math
import os
import sys
import uuid

from elasticsearch import Elasticsearch, helpers

url = os.environ.get("ES_URL")
if not url:
    print("未设置 ES_URL，跳过")
    sys.exit(0)

es = Elasticsearch(url, request_timeout=30)
suffix = uuid.uuid4().hex[:6]
alias = f"kb-{suffix}"
DIM = 64

DOCS = [
    {"id": "d1", "tenant": "t1", "title": "退货政策", "body": "签收后七天内可以申请无理由退货，商品需保持完好。"},
    {"id": "d2", "tenant": "t1", "title": "退款到账时间", "body": "退款审核通过后，原路退回，一般一到三个工作日到账。"},
    {"id": "d3", "tenant": "t1", "title": "发票开具", "body": "订单完成后可以在订单详情页申请电子发票。"},
    {"id": "d4", "tenant": "t1", "title": "会员积分", "body": "每消费一元累计一个积分，积分可以抵扣运费。"},
    {"id": "d5", "tenant": "t2", "title": "退货政策（其他租户）", "body": "本店商品不支持无理由退货。"},
    {"id": "d6", "tenant": "t1", "title": "错误码 E1024", "body": "E1024 表示支付渠道超时，请稍后重试或更换支付方式。"},
]


def embed(text: str) -> list[float]:
    """玩具嵌入：字符二元组哈希到 64 维再归一化。真实项目换成嵌入模型。"""
    v = [0.0] * DIM
    for a, b in zip(text, text[1:]):
        h = int(hashlib.md5((a + b).encode()).hexdigest(), 16)
        v[h % DIM] += 1.0
    n = math.sqrt(sum(x * x for x in v)) or 1.0
    return [x / n for x in v]


def create_index(name: str) -> None:
    es.indices.create(index=name, mappings={
        "properties": {
            "tenant": {"type": "keyword"},                       # 过滤用，不分词
            "title": {"type": "text", "analyzer": "cjk"},        # 内置 cjk 分析器（中文二元切分）；生产常用 IK 插件，以插件文档为准
            "body": {"type": "text", "analyzer": "cjk"},
            "vec": {"type": "dense_vector", "dims": DIM, "index": True, "similarity": "cosine"},
        }
    })


def load(index: str) -> None:
    # 用业务 id 作为 _id：重复同步是覆盖而不是新增（幂等）
    actions = ({"_index": index, "_id": d["id"], "_source": {**{k: v for k, v in d.items() if k != "id"}, "vec": embed(d["title"] + d["body"])}} for d in DOCS)
    ok, errors = helpers.bulk(es, actions, refresh="wait_for")
    assert ok == len(DOCS) and not errors


def bm25(q: str, tenant: str, k: int = 3) -> list[str]:
    r = es.search(index=alias, size=k, query={"bool": {
        "must": {"multi_match": {"query": q, "fields": ["title^2", "body"]}},     # 标题权重更高
        "filter": {"term": {"tenant": tenant}},                                    # 租户过滤放 filter：不算分、可缓存
    }})
    return [h["_id"] for h in r["hits"]["hits"]]


def knn(q: str, tenant: str, k: int = 3) -> list[str]:
    r = es.search(index=alias, size=k, knn={
        "field": "vec", "query_vector": embed(q), "k": k, "num_candidates": 50,
        "filter": {"term": {"tenant": tenant}},                                    # 过滤在向量检索过程中生效，而不是检索后再过滤
    })
    return [h["_id"] for h in r["hits"]["hits"]]


def rrf(*rankings: list[str], k: int = 60) -> list[str]:
    """倒数排名融合：只看名次，不看分数（BM25 分数和向量相似度不可直接比较）。"""
    score: dict[str, float] = {}
    for ranking in rankings:
        for rank, doc_id in enumerate(ranking, start=1):
            score[doc_id] = score.get(doc_id, 0.0) + 1.0 / (k + rank)
    return sorted(score, key=lambda d: (-score[d], d))


try:
    v1 = f"{alias}-v1"
    create_index(v1)
    load(v1)
    es.indices.put_alias(index=v1, name=alias)            # 业务只访问别名

    exact = bm25("E1024", "t1")
    assert exact[0] == "d6", exact
    print(f"1. 关键词检索擅长精确词（错误码）：E1024 → {exact}")

    kw = bm25("退货", "t1")
    assert "d5" not in kw and kw[0] == "d1", kw
    print(f"2. 租户过滤生效：t1 搜“退货” → {kw}（没有 t2 的 d5）")

    sem = knn("商品不想要了怎么退", "t1")
    print(f"3. 向量检索（玩具嵌入）：{sem}")
    assert "d5" not in sem

    fused = rrf(bm25("退货 退款多久到账", "t1", 5), knn("退货 退款多久到账", "t1", 5))
    assert set(fused[:2]) == {"d1", "d2"}, fused
    print(f"4. RRF 混合检索：{fused[:3]}")

    # 服务端 RRF（retriever 语法）。许可证和可用版本以官方文档为准，不可用时退回客户端融合
    try:
        r = es.search(index=alias, size=3, retriever={"rrf": {"retrievers": [
            {"standard": {"query": {"bool": {"must": {"match": {"body": "退款"}}, "filter": {"term": {"tenant": "t1"}}}}}},
            {"knn": {"field": "vec", "query_vector": embed("退款到账"), "k": 3, "num_candidates": 50, "filter": {"term": {"tenant": "t1"}}}},
        ]}})
        print(f"5. 服务端 RRF：{[h['_id'] for h in r['hits']['hits']]}")
    except Exception as e:  # noqa: BLE001  —— 例如 basic 许可下返回 403 授权错误
        print(f"5. 服务端 RRF 不可用（{type(e).__name__}: {str(e)[:80]}），使用客户端 RRF")

    # 别名切换：新建 v2（例如修改了映射），灌数据，原子切换别名，业务无感知
    v2 = f"{alias}-v2"
    create_index(v2)
    load(v2)
    es.indices.update_aliases(actions=[
        {"remove": {"index": v1, "alias": alias}},
        {"add": {"index": v2, "alias": alias}},
    ])
    assert list(es.indices.get_alias(name=alias).keys()) == [v2]
    assert bm25("E1024", "t1")[0] == "d6"
    print("6. 别名原子切换到 v2，查询不中断")
    print("全部检查通过")
finally:
    # 生产集群通常禁止通配符删除（action.destructive_requires_name），这里逐个写明索引名
    es.indices.delete(index=[f"{alias}-v1", f"{alias}-v2"], ignore_unavailable=True)
    es.close()