# -*- coding:utf-8 -*-
"""P0 探测 v2：VDB(Milvus) /v1/query 正确协议 + expr 标量过滤能力。

uapi_seed.sql 实锤协议：kNN = POST /v1/query {colname, vector, pagerows, output_fields}；
delete = POST /v1/delete {colname, pks}；expr 参数名待实测（v1 探测报错实锤存在 expr 概念）。
"""
import asyncio
import json
import os
import sys
import time

WORKDIR = "/d/pipeline/pipeline-app"
os.chdir(WORKDIR)
sys.path.insert(0, WORKDIR)

from appPublic.folderUtils import ProgramPath          # noqa: E402
from appPublic.jsonConfig import getConfig             # noqa: E402
from appPublic.event_dispatcher import EventDispatcher # noqa: E402
from sqlor.dbpools import DBPools                      # noqa: E402
from ahserver.serverenv import ServerEnv               # noqa: E402

p = ProgramPath()
config = getConfig(WORKDIR, NS={'workdir': WORKDIR, 'ProgramPath': p})
DBPools(config.databases)
se = ServerEnv()
se.event_dispatcher = EventDispatcher()
se.get_module_dbname = lambda m: 'pipeline'

TEST_COL = "opp_p0_probe2_%d" % int(time.time())
DIM = 1024


async def _sql(sql, args=None):
    async with DBPools().sqlorContext("pipeline") as sor:
        recs = await sor.sqlExe(sql, args or {})
    return recs or []


def vec(seed):
    v = [0.0] * DIM
    v[seed] = 1.0
    v[seed + 1] = 0.1
    return v


async def main():
    out = {"steps": [], "col": TEST_COL}

    def step(name, ok, detail=""):
        out["steps"].append({"name": name, "ok": bool(ok), "detail": str(detail)[:600]})
        print("[%s] %s %s" % ("OK " if ok else "FAIL", name, str(detail)[:400]))

    recs = await _sql("SELECT baseurl FROM upapp WHERE id='rag-vdb'")
    vdb_base = (recs[0].baseurl or "").rstrip("/") if recs else ""
    if not vdb_base:
        step("vdb_base", False, "未配置")
        return

    import aiohttp
    timeout = aiohttp.ClientTimeout(total=30)
    async with aiohttp.ClientSession(timeout=timeout) as s:
        async def post(path, payload):
            r = await s.post(vdb_base + path, json=payload)
            return r.status, await r.text()

        st, body = await post("/v1/createcollection", {"colname": TEST_COL, "fields": [
            {"name": "id", "type": "str", "is_primary": True, "max_length": 64},
            {"name": "vector", "type": "fvector", "dim": DIM},
            {"name": "text", "type": "str", "max_length": 65535},
            {"name": "batch_id", "type": "str", "max_length": 64}],
            "description": "P0 probe2", "metric": "COSINE"})
        step("createcollection(+batch_id字段)", st == 200 and "SUCCEEDED" in body, body)

        st, body = await post("/v1/upsert", {"colname": TEST_COL, "data": [
            {"id": "t1", "vector": vec(0), "text": "外包人员管理系统", "batch_id": "b1"},
            {"id": "t2", "vector": vec(2), "text": "合同管理系统", "batch_id": "b1"},
            {"id": "t3", "vector": vec(4), "text": "知识库问答", "batch_id": "b2"}]})
        step("upsert(3)", st == 200 and "SUCCEEDED" in body, body)

        # kNN 检索（生产协议：/v1/query + pagerows）
        st, body = await post("/v1/query", {
            "colname": TEST_COL, "vector": vec(0), "pagerows": 2,
            "output_fields": ["id", "text", "batch_id"]})
        ok = st == 200 and "t1" in body
        step("query.kNN(pagerows=2)", ok, body)

        # 标量过滤 expr（升级路径①的关键：常驻 collection + batch_id filter）
        for key in ("expr", "filter"):
            st, body = await post("/v1/query", {
                "colname": TEST_COL, key: 'batch_id == "b1"',
                "output_fields": ["id", "text"]})
            ok = st == 200 and "t1" in body and "t3" not in body
            step("query.scalar(%s=...)" % key, ok, body)

        # kNN + expr 组合（聚类工作集免拷贝的终极验证）
        st, body = await post("/v1/query", {
            "colname": TEST_COL, "vector": vec(0), "pagerows": 3,
            "expr": 'batch_id == "b1"', "output_fields": ["id", "batch_id"]})
        ok = st == 200 and "t1" in body and "t3" not in body
        step("query.kNN+expr", ok, body)

        # 按主键删除（增量缓存维护用）
        st, body = await post("/v1/delete", {"colname": TEST_COL, "pks": ["t3"]})
        step("delete(pks)", st == 200 and "SUCCEEDED" in body, body)
        st, body = await post("/v1/query", {
            "colname": TEST_COL, "vector": vec(4), "pagerows": 2,
            "output_fields": ["id"]})
        step("delete生效验证(t3应不在)", st == 200 and "t3" not in body, body)

        # 清理
        st, body = await post("/v1/dropcollection", {"colname": TEST_COL})
        step("dropcollection(清理)", st == 200 and "SUCCEEDED" in body, body)

    print("=== RESULT_JSON ===")
    print(json.dumps(out, ensure_ascii=False))


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except Exception:
        import traceback
        traceback.print_exc()
        sys.exit(1)