#!/usr/bin/env python3
import json, time, uuid, sys
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:
        # try nested
        for k, v in data.items():
            if isinstance(v, dict):
                token = v.get("access_token") or v.get("token")
                if token:
                    break
    if not token:
        print("LOGIN_PAYLOAD:", json.dumps(data)[:500])
        sys.exit(1)
    return token

class WSClient:
    def __init__(self, token):
        url = f"ws://10.168.3.36:8081/api/user/friend-agent/ws?access_token={token}"
        self.ws = websocket.create_connection(url, timeout=90)
        self.pending = {}
        self.thread = None
    def _read_loop(self):
        while True:
            try:
                msg = self.ws.recv()
            except Exception:
                return
            try:
                obj = json.loads(msg)
            except Exception:
                continue
            rid = obj.get("request_id")
            if rid and rid in self.pending:
                self.pending.pop(rid).put(obj)
    def request(self, method, params=None, timeout=60):
        import threading
        if self.thread is None:
            self.thread = threading.Thread(target=self._read_loop, daemon=True)
            self.thread.start()
        rid = f"q-{uuid.uuid4().hex[:8]}"
        import queue
        q = queue.Queue()
        self.pending[rid] = q
        self.ws.send(json.dumps({"type": "req", "request_id": rid, "method": method, "params": params or {}}))
        try:
            return q.get(timeout=timeout)
        except queue.Empty:
            self.pending.pop(rid, None)
            return {"status": "error", "error": {"message": "timeout"}}

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

def main():
    token = login()
    print("login ok")
    c = WSClient(token)
    time.sleep(0.5)
    results = {}

    # 1. llm.model.probe with model_id
    resp = c.request("llm.model.probe", {"model_id": 1})
    results["llm.model.probe(model_id:1)"] = check("llm.model.probe {model_id:1}", resp, want_ok=True)

    # 2. llm.model.clear-cooldown
    resp = c.request("llm.model.clear-cooldown", {"model_id": 1})
    results["llm.model.clear-cooldown(model_id:1)"] = check("llm.model.clear-cooldown {model_id:1}", resp, want_ok=True)

    # 3. create test channel
    resp = c.request("llm.channels.create", {"name": "ws-test-channel", "protocol": "openai", "api_base": "http://127.0.0.1:59999/v1", "api_key": "sk-test", "enabled": True})
    ok = check("llm.channels.create", resp, want_ok=True)
    results["llm.channels.create"] = ok
    new_id = None
    if ok:
        res = resp.get("result")
        if isinstance(res, dict):
            new_id = res.get("id") or res.get("channel_id")
    print("  new channel id:", new_id)

    # 4. llm.channel.test on channel 1
    resp = c.request("llm.channel.test", {"id": 1}, timeout=90)
    results["llm.channel.test(id:1)"] = check("llm.channel.test {id:1}", resp, want_ok=True)

    # 5. llm.channel.sync-models on channel 1
    resp = c.request("llm.channel.sync-models", {"id": 1}, timeout=90)
    results["llm.channel.sync-models(id:1)"] = check("llm.channel.sync-models {id:1}", resp, want_ok=True)

    # 6. model CRUD on test channel
    if new_id:
        resp = c.request("llm.models.create", {"channel_id": new_id, "model_key": "ws-test-model", "upstream_model_id": "ws-test-model"})
        mok = check("llm.models.create", resp, want_ok=True)
        results["llm.models.create"] = mok
        mid = None
        if mok:
            res = resp.get("result")
            if isinstance(res, dict):
                mid = res.get("id")
        print("  new model id:", mid)
        if mid:
            resp = c.request("llm.models.update", {"id": mid, "model_key": "ws-test-model", "upstream_model_id": "ws-test-model", "display_name": "WS Test Model"})
            results["llm.models.update"] = check("llm.models.update", resp, want_ok=True)
            resp = c.request("llm.models.delete", {"id": mid})
            results["llm.models.delete"] = check("llm.models.delete", resp, want_ok=True)
        # delete test channel
        resp = c.request("llm.channels.delete", {"id": new_id})
        results["llm.channels.delete"] = check("llm.channels.delete", resp, want_ok=True)
    else:
        print("  SKIP model CRUD (no channel id)")

    # 7. llm.channels.update on test channel (if created)
    if new_id:
        pass  # already deleted; update tested implicitly via create flow

    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")
    c.ws.close()

if __name__ == "__main__":
    main()