#!/usr/bin/env python3
"""
probe_cliproxy.py — Thăm dò CLIProxy để tìm command/channel hợp lệ
"""
import socket
import sys

HOST = "127.0.0.1"
PORT = 8317
PASSWORD = "123456"


# ── RESP helpers ────────────────────────────────────────────────────────────

def resp_command(*parts):
    data = [f"*{len(parts)}\r\n".encode()]
    for part in parts:
        b = str(part).encode()
        data.append(f"${len(b)}\r\n".encode())
        data.append(b + b"\r\n")
    return b"".join(data)


class Client:
    def __init__(self, host, port, password, timeout=5):
        self.sock = socket.create_connection((host, port), timeout=timeout)
        self.file = self.sock.makefile("rb")
        self._send("AUTH", password)
        reply = self._read()
        if not (isinstance(reply, str) and reply.upper() == "OK"):
            raise RuntimeError(f"AUTH failed: {reply!r}")

    def close(self):
        try: self.file.close()
        finally: self.sock.close()

    def _send(self, *parts):
        self.sock.sendall(resp_command(*parts))

    def _read_line(self):
        line = self.file.readline()
        if not line:
            raise EOFError("Connection closed")
        return line.rstrip(b"\r\n")

    def _read(self):
        line = self._read_line()
        prefix, payload = line[:1], line[1:]
        if prefix == b"+": return payload.decode()
        if prefix == b"-": raise RuntimeError(payload.decode())
        if prefix == b":": return int(payload)
        if prefix == b"$":
            length = int(payload)
            if length == -1: return None
            data = self.file.read(length)
            self.file.read(2)
            return data.decode("utf-8", "replace")
        if prefix == b"*":
            count = int(payload)
            if count == -1: return None
            return [self._read() for _ in range(count)]
        raise RuntimeError(f"Unknown RESP prefix: {line!r}")

    def cmd(self, *parts):
        """Gửi command, trả về (ok, result). Không raise exception."""
        self._send(*parts)
        try:
            return True, self._read()
        except RuntimeError as e:
            return False, str(e)
        except Exception as e:
            return False, f"[{type(e).__name__}] {e}"


# ── Probe logic ──────────────────────────────────────────────────────────────

def section(title):
    print(f"\n{'─'*55}")
    print(f"  {title}")
    print(f"{'─'*55}")


def probe(c):
    # 1. Các meta-command thường gặp
    section("1. Meta commands (HELP / INFO / PING / COMMAND)")
    for cmd in ["PING", "HELP", "INFO", "COMMAND", "COMMAND COUNT", "COMMAND DOCS"]:
        ok, val = c.cmd(*cmd.split())
        status = "✅ OK" if ok else "❌"
        preview = str(val)[:120].replace("\r\n", "\\r\\n") if val else "nil"
        print(f"  {cmd:<20} {status}  →  {preview}")

    # 2. Thử RPOP với các channel name phổ biến
    section("2. RPOP — thử các channel name")
    channels = [
        "events", "queue", "logs", "usage", "requests",
        "stream", "data", "codex", "proxy", "api",
        "messages", "output", "items", "records",
    ]
    valid_channels = []
    for ch in channels:
        ok, val = c.cmd("RPOP", ch)
        if ok:
            preview = str(val)[:80] if val else "nil (rỗng)"
            print(f"  RPOP {ch:<16} ✅ OK  →  {preview}")
            valid_channels.append(ch)
        else:
            print(f"  RPOP {ch:<16} ❌  {val}")

    # 3. Thử LRANGE (nếu server dùng list-based)
    section("3. LRANGE — xem nội dung list (không xoá)")
    for ch in channels:
        ok, val = c.cmd("LRANGE", ch, "0", "2")
        if ok:
            preview = str(val)[:100] if val else "[]"
            print(f"  LRANGE {ch:<14} ✅ OK  →  {preview}")

    # 4. Thử LLEN (độ dài list)
    section("4. LLEN — độ dài list")
    for ch in channels:
        ok, val = c.cmd("LLEN", ch)
        if ok:
            print(f"  LLEN {ch:<16} ✅ OK  →  {val}")

    # 5. Thử KEYS nếu server hỗ trợ
    section("5. KEYS * — liệt kê tất cả key")
    ok, val = c.cmd("KEYS", "*")
    if ok:
        print(f"  ✅ KEYS * →  {val}")
    else:
        print(f"  ❌ KEYS * →  {val}")

    # 6. Thử SUBSCRIBE/PSUBSCRIBE (pub/sub mode)
    section("6. SUBSCRIBE — thử pub/sub channels")
    for ch in ["events", "logs", "usage", "requests"]:
        ok, val = c.cmd("SUBSCRIBE", ch)
        if ok:
            print(f"  SUBSCRIBE {ch:<12} ✅ OK  →  {val}")
        else:
            print(f"  SUBSCRIBE {ch:<12} ❌  {val}")

    # 7. Thử XREAD / XLEN (Redis Streams)
    section("7. XREAD / XLEN — thử Redis Streams")
    for ch in ["events", "logs", "usage"]:
        ok, val = c.cmd("XLEN", ch)
        if ok:
            print(f"  XLEN {ch:<16} ✅ OK  →  {val}")
        else:
            print(f"  XLEN {ch:<16} ❌  {val}")

    # 8. Tổng kết
    section("TỔNG KẾT")
    if valid_channels:
        print(f"  ✅ Channel RPOP hợp lệ: {valid_channels}")
        print(f"\n  → Thêm vào config.json:")
        print(f'     "queue_channel": "{valid_channels[0]}"')
    else:
        print("  ⚠️  Không tìm thấy channel RPOP hợp lệ nào.")
        print("  → CLIProxy có thể dùng giao thức khác (pub/sub, stream, HTTP).")
        print("  → Kiểm tra tài liệu hoặc source code của CLIProxy.")


# ── Main ─────────────────────────────────────────────────────────────────────

def main():
    print(f"Kết nối tới CLIProxy {HOST}:{PORT} ...")
    try:
        c = Client(HOST, PORT, PASSWORD)
    except Exception as e:
        print(f"❌ Kết nối thất bại: {e}", file=sys.stderr)
        sys.exit(1)
    print(f"✅ AUTH OK\n")
    try:
        probe(c)
    finally:
        c.close()


if __name__ == "__main__":
    main()