"""Smoke-test WebSocket RAG streaming against running service."""

import asyncio
import json
import sys
import urllib.error
import urllib.request

import websockets


async def main() -> int:
    login_body = json.dumps({"username": "admin", "password": "admin123"}).encode()
    login_req = urllib.request.Request(
        "http://127.0.0.1:8000/api/auth/login",
        data=login_body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(login_req, timeout=10) as resp:
            token = json.loads(resp.read().decode())["access_token"]
    except urllib.error.HTTPError as e:
        print(f"Login failed: {e.read().decode()}", file=sys.stderr)
        return 1

    uri = f"ws://127.0.0.1:8000/ws?token={token}"
    types: list[str] = []
    chunks = 0

    async with websockets.connect(uri, open_timeout=10) as ws:
        await ws.send(json.dumps({
            "action": "rag_query_global",
            "question": "Кратко: что есть в базе знаний?",
            "history": [],
            "chat_mode": "hybrid",
            "retrieval_mode": "naive",
        }))

        while len(types) < 200:
            raw = await asyncio.wait_for(ws.recv(), timeout=180)
            msg = json.loads(raw)
            msg_type = msg.get("type")
            types.append(msg_type)
            if msg_type == "rag_chunk":
                chunks += 1
            if msg_type in ("rag_response", "rag_error"):
                if msg_type == "rag_error":
                    print(f"RAG error: {msg.get('error')}", file=sys.stderr)
                    return 1
                print(f"OK: types={types[:5]}... chunks={chunks} answer_len={len(msg.get('answer', ''))}")
                return 0

    print("No terminal message received", file=sys.stderr)
    return 1


if __name__ == "__main__":
    raise SystemExit(asyncio.run(main()))