#!/usr/bin/env python3
import json, time, uuid, sys, threading, queue
import requests
import websocket

BASE = "http://10.168.3.36:8081"

def login():
    r = requests.post(f"{BASE}/api/login", json={"username": "cxbsoft", "password": "cf5584aca177a811ffda8f4789e2e143"}, timeout=10)
    r.raise_for_status()
    data = r.json()
    token = data.get("access_token") or data.get("data", {}).get("access_token") or data.get("token")
    if not token:
        for k, v in data.items():
            if isinstance(v, dict):
                token = v.get("access_token") or v.get("token")
                if token:
                    break
    return token

def main():
    token = login()
    url = f"ws://10.168.3.36:8081/api/user/friend-agent/ws?access_token={token}"
    ws = websocket.create_connection(url, timeout=30)
    pending = {}
    stop = False
    def read_loop():
        while not stop:
            try:
                msg = ws.recv()
            except Exception:
                return
            try:
                obj = json.loads(msg)
            except Exception:
                continue
            rid = obj.get("request_id")
            if rid and rid in pending:
                pending.pop(rid).put(obj)
    t = threading.Thread(target=read_loop, daemon=True)
    t.start()
    time.sleep(1)
    def call(method, params=None, timeout=60):
        rid = f"q-{uuid.uuid4().hex[:8]}"
        q = queue.Queue()
        pending[rid] = q
        ws.send(json.dumps({"type": "req", "request_id": rid, "method": method, "params": params or {}}))
        try:
            return q.get(timeout=timeout)
        except queue.Empty:
            return {"status": "error", "error": {"message": "timeout"}}

    results = {}
    def check(name, resp, want_ok=True):
        status = resp.get("status")
        ok = (status == "ok") if want_ok else (status != "ok")
        results[name] = ok
        body = json.dumps(resp, ensure_ascii=False)
        if len(body) > 300:
            body = body[:300] + "..."
        print(f"[{'PASS' if ok else 'FAIL'}] {name}: {body}")
        return ok

    # P1: 14 subjects
    check("llm.health", call("llm.health"))
    check("llm.health.reset", call("llm.health.reset"))
    check("llm.health.events{limit:20}", call("llm.health.events", {"limit": 20}))
    check("llm.overview", call("llm.overview"))
    check("llm.logs{limit:10}", call("llm.logs", {"limit": 10}))
    check("llm.dashboard", call("llm.dashboard", {"days": 7}))
    check("llm.channel.stats{id:1}", call("llm.channel.stats", {"id": 1}))
    check("llm.groups.list", call("llm.groups.list"))

    # groups CRUD
    resp = call("llm.groups.create", {"name": "p1-ws-test-group", "description": "p1 ws test"})
    gok = check("llm.groups.create", resp)
    gid = None
    if gok:
        res = resp.get("result")
        if isinstance(res, dict):
            gid = res.get("id")
    print("  group id:", gid)
    if gid:
        check("llm.groups.update", call("llm.groups.update", {"id": gid, "name": "p1-ws-test-group", "description": "p1 ws test updated"}))
        check("llm.groups.delete", call("llm.groups.delete", {"id": gid}))

    # regression: P0 + existing subjects
    check("REG llm.channels.list", call("llm.channels.list"))
    check("REG llm.model.probe{id:9}", call("llm.model.probe", {"id": 9}, timeout=90))
    check("REG llm.channel.create", call("llm.channels.create", {"name": "p1-reg-ch", "protocol": "openai", "api_base": "http://127.0.0.1:59998/v1"}))
    reg_id = None
    resp2 = pending.get("last")
    # get the created channel id from last response
    # re-fetch channels list to find it
    resp3 = call("llm.channels.list")
    if resp3.get("status") == "ok":
        chs = resp3.get("result", {}).get("channels", [])
        for c in chs:
            if c.get("name") == "p1-reg-ch":
                reg_id = c.get("id")
                break
    if reg_id:
        check("REG llm.channel.delete", call("llm.channels.delete", {"id": reg_id}))

    print("\n=== SUMMARY ===")
    all_ok = True
    for k, v in results.items():
        print(f"  {'OK ' if v else 'BAD'} {k}")
        all_ok = all_ok and v
    print("ALL_GREEN" if all_ok else "SOME_FAILED")
    stop = True
    ws.close()

if __name__ == "__main__":
    main()