#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""专题十四 · 步骤 5：命中测试计时（走知识库当前检索配置）。"""
from __future__ import annotations

import sys
import time

import _bootstrap  # noqa: F401
import _config as C
from dify_api import Dify


def main() -> int:
    d = Dify(C.DIFY_EMAIL, C.DIFY_PASSWORD)
    try:
        d.login()
        t0 = time.time()
        r = d.post(
            "/datasets/%s/hit-testing" % C.DS_ID,
            body={"query": C.PROBE_QUERY_KB},
            extra="-m 180 --max-time 180",
            t=200,
        )
        elapsed = round(time.time() - t0, 2)
        print("HTTP", r.get("code"), "elapsed", elapsed, "s")
        j = r.get("json")
        if isinstance(j, list):
            recs = j
        else:
            recs = (j or {}).get("records") or []
        for i, rec in enumerate(recs[:5], 1):
            seg = rec.get("segment") or rec
            print(i, rec.get("score"), (seg.get("content") or "")[:80])
        return 0 if r.get("code") == "200" else 1
    finally:
        d.close()


if __name__ == "__main__":
    sys.exit(main())