from __future__ import annotations

import json
import math
import time
from typing import Any
from urllib.error import HTTPError
from urllib.request import Request, urlopen

from pymilvus import DataType, Function, FunctionType, MilvusClient

# ==================== 连接配置 ====================
MILVUS_HOST = "http://c-xxx.milvus.aliyuncs.com"
MILVUS_TOKEN = "root:xxx"
# RESTful 接口与 gRPC 同在 19530 端口，需显式带端口（否则默认 80 端口会连接超时）
MILVUS_REST_BASE_URL = f"{MILVUS_HOST}:19530"

client = MilvusClient(uri=MILVUS_HOST, token=MILVUS_TOKEN)
collection_name = "shop_reviews"


def post_json(path: str, body: dict[str, Any], timeout: int = 120, retries: int = 3) -> tuple[int, dict[str, Any]]:
    """AI_RERANK / AI_SENTIMENT / AI_SUMMARIZE 的 REST 即时接口统一封装（走大模型，加重试更稳）。"""
    last = None
    for _ in range(retries):
        request = Request(
            f"{MILVUS_REST_BASE_URL.rstrip('/')}{path}",
            data=json.dumps(body, ensure_ascii=False).encode("utf-8"),
            headers={"Authorization": f"Bearer {MILVUS_TOKEN}", "Content-Type": "application/json"},
            method="POST",
        )
        try:
            with urlopen(request, timeout=timeout) as response:
                status, data = response.status, json.loads(response.read().decode("utf-8"))
        except HTTPError as exc:
            status, data = exc.code, json.loads(exc.read().decode("utf-8"))
        last = (status, data)
        if status == 200 and data.get("code") == 0:
            return status, data
        time.sleep(1)
    return last


def circle_to_wkt(lng: float, lat: float, radius_km: float, num_seg: int = 32) -> str:
    """把以 (lng,lat) 为圆心、radius_km 为半径的圆，近似成正多边形的 WKT（用于半径粗过滤）。
    注意：经纬度是角度，此处按平面近似换算，纬度越高误差越大。"""
    d_lat = radius_km / 111.0
    d_lng = radius_km / (111.0 * math.cos(math.radians(lat)))

    pts = [ ]

    for i in range(num_seg + 1):
        theta = 2 * math.pi * i / num_seg
        pts.append(f"{lng + d_lng * math.cos(theta):.6f} {lat + d_lat * math.sin(theta):.6f}")
    return f"POLYGON(({', '.join(pts)}))"


# ==================== 步骤 1：建带地理字段 + embedding Function 的 Collection ====================
# location 用 DataType.GEOMETRY 存 POI 点；review 挂 AI_EMBEDDING，写入时自动生成 1024 维向量
if client.has_collection(collection_name):
    client.drop_collection(collection_name)

schema = MilvusClient.create_schema(auto_id=True, enable_dynamic_field=False)
schema.add_field("id", DataType.INT64, is_primary=True)
schema.add_field("shop_id", DataType.INT64)
schema.add_field("shop_name", DataType.VARCHAR, max_length=256)
schema.add_field("category", DataType.VARCHAR, max_length=64)
schema.add_field("location", DataType.GEOMETRY)                       # POI 经纬度点（WKT: POINT(lng lat)）
schema.add_field("review", DataType.VARCHAR, max_length=4096)         # 评价原文，作为 embedding 输入
schema.add_field("review_embedding", DataType.FLOAT_VECTOR, dim=1024)  # 评价向量，维度需与 dim 一致

# 写入时自动把 review 向量化
schema.add_function(
    Function(
        name="embed_review",
        function_type=FunctionType.TEXTEMBEDDING,
        input_field_names=["review"],
        output_field_names=["review_embedding"],
        params={
            "provider": "aliyun_milvus",
            "model_name": "qwen3.7-text-embedding",
            "dim": 1024,
            "max_client_batch_size": 20,   # qwen3.7-text-embedding 百炼官方文档批量大小为 20
            "max_concurrency": 1,
        },
    )
)

index_params = client.prepare_index_params()
index_params.add_index(
    field_name="review_embedding",
    index_type="HNSW",
    metric_type="COSINE",
    params={"M": 16, "efConstruction": 200},
)
client.create_collection(
    collection_name=collection_name,
    schema=schema,
    index_params=index_params,
)


# ==================== 步骤 2：POI 与评价数据入库（review_embedding 由 Function 自动生成，无需手工传向量） ====================
rows = [
    {"shop_id": 1001, "shop_name": "巴蜀人家·亲子川菜", "category": "川菜",
     "location": "POINT(116.4612 39.9088)",
     "review": "带娃来的，有儿童座椅和小碗餐具，毛血旺不算太辣，服务员很有耐心，环境干净。"},
    {"shop_id": 1001, "shop_name": "巴蜀人家·亲子川菜", "category": "川菜",
     "location": "POINT(116.4612 39.9088)",
     "review": "周末人多要等位，但菜品稳定，适合一家人聚餐，孩子吃得开心。"},
    {"shop_id": 1002, "shop_name": "麻辣江湖", "category": "川菜",
     "location": "POINT(116.4501 39.9101)",
     "review": "重口味爱好者天堂，水煮鱼够味，但空间偏小、比较吵，不太适合带小孩。"},
    {"shop_id": 1003, "shop_name": "静巷咖啡", "category": "咖啡",
     "location": "POINT(116.4550 39.9075)",
     "review": "环境安静适合谈事，手冲不错，不过跟川菜需求无关。"},
]
client.insert(collection_name, rows)
client.flush(collection_name)
client.load_collection(collection_name)


# ==================== 步骤 3：地理过滤查询（ST_WITHIN 商圈多边形 / 半径近似圆） ====================
# 3.1 查询某商圈多边形范围内的门店
cbd_polygon = "POLYGON((116.448 39.905, 116.465 39.905, 116.465 39.914, 116.448 39.914, 116.448 39.905))"
hits = client.query(
    collection_name=collection_name,
    filter=f"ST_WITHIN(location, '{cbd_polygon}')",
    output_fields=["shop_name", "category", "location"],
    limit=100,
)
print("\n" + "=" * 64)
print("步骤 3 | 地理过滤查询  ST_WITHIN(location, 多边形 / 半径圆)")
print("=" * 64)
print("[3.1] CBD 商圈多边形范围内的门店：")
for h in hits:
    print(f"      · {h['shop_name']}（{h['category']}）  {h['location']}")

# 3.2 附近 3km 半径查询：把圆近似成多边形，再用 ST_WITHIN 判断 POI 是否落在圈内
center_lng, center_lat = 116.4600, 39.9090
circle = circle_to_wkt(center_lng, center_lat, radius_km=3.0)
nearby = client.query(
    collection_name=collection_name,
    filter=f"ST_WITHIN(location, '{circle}')",
    output_fields=["shop_name", "category", "location"],
    limit=200,
)
print(f"[3.2] 中心点 ({center_lng}, {center_lat}) 附近 3km 半径内的门店：")
for n in nearby:
    print(f"      · {n['shop_name']}（{n['category']}）  {n['location']}")


# ==================== 步骤 4：语义检索 + 地理 filter 融合进一次 search() ====================
# ANN 走评价向量负责「对味」，filter 里的 ST_ 空间函数负责「就近」，一次请求联合求解
QUERY = "附近适合亲子的川菜馆"
geo_filter = f"category == '川菜' && ST_WITHIN(location, '{circle}')"
results = client.search(
    collection_name=collection_name,
    data=[QUERY],
    anns_field="review_embedding",
    filter=geo_filter,
    limit=20,
    output_fields=["shop_id", "shop_name", "category", "location", "review"],
)
print("\n" + "=" * 64)
print("步骤 4 | 语义检索 + 地理过滤 融合查询（一次 search 联合求解）")
print("=" * 64)
print(f"查询语句：{QUERY}")
print("过滤条件：category == '川菜'  且落在附近 3km 圈内")
print("-" * 64)
for rank, hit in enumerate(results[0], 1):
    e = hit["entity"]
    print(f"  {rank}. [相似度 {hit['distance']:.4f}]  {e['shop_name']} ｜ {e['review'][:24]}")


# ==================== 步骤 5：AI_RERANK 重排候选 ====================
RERANK_MODEL = "qwen3-rerank"

# 5.1 在 search() 中挂 RERANK ranker（reranker="model"），ANN 召回后由模型二次精排
reranker = Function(
    name="rerank_reviews",
    function_type=FunctionType.RERANK,
    input_field_names=["review"],
    params={
        "reranker": "model",
        "provider": "aliyun_milvus",
        "model_name": RERANK_MODEL,
        "queries": [QUERY],
        "timeout_sec": 10,
    },
)
results = client.search(
    collection_name=collection_name,
    data=[QUERY],
    anns_field="review_embedding",
    filter=geo_filter,
    limit=20,
    output_fields=["shop_id", "shop_name", "review"],
    ranker=reranker,
)
print("\n" + "=" * 64)
print("步骤 5 | AI_RERANK 大模型重排（qwen3-rerank）")
print("=" * 64)
print("[5.1] search() 内置 ranker：ANN 召回后由大模型二次精排")
print("-" * 64)
for rank, hit in enumerate(results[0], 1):
    e = hit["entity"]
    print(f"  {rank}. [重排分 {hit['distance']:.4f}]  {e['shop_name']} ｜ {e['review'][:24]}")

# 5.2 也可用无状态 REST 即时接口 /v2/vectordb/ai/rerank
#     （documents 必须是 string 数组、query 为顶层 string、不支持 top_n，每个 doc 返回一个分）
rerank_docs = [
    "带娃来的，有儿童座椅和小碗餐具，服务员很有耐心。",
    "重口味爱好者天堂，水煮鱼够味，但空间偏小、比较吵，不太适合带小孩。",
]
status, data = post_json(
    "/v2/vectordb/ai/rerank",
    {
        "model_name": RERANK_MODEL,
        "query": QUERY,
        "documents": rerank_docs,
        "params": {"timeout_sec": 10},
    },
)
assert status == 200 and data.get("code") == 0, data
print(f"\n[5.2] REST 即时接口 /v2/vectordb/ai/rerank（query：{QUERY}）")
print("-" * 64)
ranked = sorted(data["data"]["output"]["results"], key=lambda x: x["relevance_score"], reverse=True)
for rank, item in enumerate(ranked, 1):
    print(f"  {rank}. [相关度 {item['relevance_score']:.4f}]  {rerank_docs[item['index']][:28]}")


# ==================== 步骤 6：AI_SENTIMENT 算好评率 + AI_SUMMARIZE 生成店铺摘要 ====================
shop_id = 1001
reviews = [r["review"] for r in client.query(
    collection_name=collection_name,
    filter=f"shop_id == {shop_id}",
    output_fields=["review"],
    limit=1000,
)]

# 6.1 情感分类 → 好评率
status, sent = post_json("/v2/vectordb/ai/sentiment", {
    "model_name": "qwen3.7-max",
    "texts": reviews,
    "params": {"categories": ["positive", "negative", "neutral"], "temperature": 0},
})
assert status == 200 and sent.get("code") == 0, sent
labels = sent["data"]["output"]["outputs"]
valid = [x for x in labels if x]                      # null 计入「未判定」，不参与聚合
positive = sum(1 for x in valid if x == "positive")
good_rate = positive / len(valid) if valid else 0.0
print("\n" + "=" * 64)
print(f"步骤 6 | AI_SENTIMENT 好评率 + AI_SUMMARIZE 摘要（shop_id={shop_id}）")
print("=" * 64)
print(f"[6.1] 好评率：{good_rate:.0%}   （有效样本 {len(valid)}/{len(labels)}）")

# 6.2 摘要 → 店铺评价一句话
status, summ = post_json("/v2/vectordb/ai/summarize", {
    "model_name": "qwen3.7-max",
    "texts": ["\n".join(reviews)],
    "params": {"max_words": 40, "temperature": 0,
               "prompt": "用中文概括这家店的整体评价，突出招牌菜、适合人群与环境。"},
})
assert status == 200 and summ.get("code") == 0, summ
print(f"[6.2] 店铺摘要：{summ['data']['output']['outputs'][0]}")