#!/usr/bin/env python3
"""RAG 追问链烟测：防止“分别是什么”回退成 count。"""

from __future__ import annotations

import os
import sys
import httpx


def main() -> int:
    default_base = f"http://localhost:{os.getenv('BACKEND_PORT', '8000')}"
    base = os.getenv("SMOKE_BASE_URL", default_base).rstrip("/")
    username = os.getenv("SMOKE_USER", "boss")
    password = os.getenv("SMOKE_PASSWORD", "boss123")

    with httpx.Client(timeout=30.0) as client:
        login = client.post(
            f"{base}/auth/login",
            headers={"Content-Type": "application/x-www-form-urlencoded"},
            data={"username": username, "password": password},
        )
        if login.status_code != 200:
            print(f"[FAIL] login status={login.status_code} body={login.text[:240]}")
            return 2

        token = login.json().get("access_token")
        if not token:
            print("[FAIL] login token missing")
            return 2

        headers = {"Authorization": f"Bearer {token}"}

        q0 = client.post(
            f"{base}/rag/ask",
            headers=headers,
            json={"question": "健身食谱项目里面有多少个信号源", "asked_via": "pmflow"},
        )
        if q0.status_code != 200:
            print(f"[FAIL] q0 status={q0.status_code} body={q0.text[:240]}")
            return 2

        q1 = client.post(
            f"{base}/rag/ask",
            headers=headers,
            json={"question": "目前系统里面有多少个项目", "asked_via": "pmflow"},
        )
        if q1.status_code != 200:
            print(f"[FAIL] q1 status={q1.status_code} body={q1.text[:240]}")
            return 2

        q2 = client.post(
            f"{base}/rag/ask",
            headers=headers,
            json={"question": "分别是什么", "asked_via": "pmflow"},
        )
        if q2.status_code != 200:
            print(f"[FAIL] q2 status={q2.status_code} body={q2.text[:240]}")
            return 2

        d0 = (q0.json() or {}).get("data") or {}
        d1 = (q1.json() or {}).get("data") or {}
        d2 = (q2.json() or {}).get("data") or {}
        qt0 = str(d0.get("query_type") or "")
        qt1 = str(d1.get("query_type") or "")
        qt2 = str(d2.get("query_type") or "")
        a0 = str(d0.get("answer") or "")
        a2 = str(d2.get("answer") or "")

        q3 = client.post(
            f"{base}/rag/ask",
            headers=headers,
            json={"question": "健身食谱项目里面有多少个信号源", "asked_via": "pmflow"},
        )
        if q3.status_code != 200:
            print(f"[FAIL] q3 status={q3.status_code} body={q3.text[:240]}")
            return 2

        q4 = client.post(
            f"{base}/rag/ask",
            headers=headers,
            json={"question": "分别什么信号源", "asked_via": "pmflow"},
        )
        if q4.status_code != 200:
            print(f"[FAIL] q4 status={q4.status_code} body={q4.text[:240]}")
            return 2

        d3 = (q3.json() or {}).get("data") or {}
        d4 = (q4.json() or {}).get("data") or {}
        qt3 = str(d3.get("query_type") or "")
        qt4 = str(d4.get("query_type") or "")
        a4 = str(d4.get("answer") or "")

        ok = (
            (qt0 == "count")
            and ("信号" in a0)
            and (qt1 == "count")
            and ("项目总数" in str(d1.get("answer") or ""))
            and (qt2 == "list")
            and ("项目列表" in a2 or "#" in a2)
            and (qt3 == "count")
            and (qt4 == "list")
            and ("信号清单" in a4 or "信号" in a4)
        )
        if not ok:
            print(
                f"[FAIL] qt0={qt0} qt1={qt1} qt2={qt2} qt3={qt3} qt4={qt4} "
                f"a0={a0[:120]} a1={str(d1.get('answer') or '')[:120]} a2={a2[:120]} a4={a4[:120]}"
            )
            return 2

        print("[OK] followup smoke passed")
        return 0


if __name__ == "__main__":
    raise SystemExit(main())