import struct
import socket

def binary_get(sock, key: str, getk: bool = False):
    """二进制 GET 命令"""
    key_bytes = key.encode()
    opcode = 0x0C if getk else 0x00

    header = struct.pack(
        "!BBHBBHIIQ",
        0x80,           # magic
        opcode,         # opcode
        len(key_bytes), # key length
        0,              # extras length
        0x00,           # data type
        0,              # reserved
        len(key_bytes), # body length
        1,              # opaque
        0               # cas
    )

    sock.sendall(header + key_bytes)

    # 读取响应头
    resp_header = _recv_exact(sock, 24)
    (_, opcode, key_len, extras_len, _, status, body_len, opaque, cas) = \
        struct.unpack("!BBHBBHIIQ", resp_header)

    # 读取响应体
    body = _recv_exact(sock, body_len)

    if status == 0x0001:  # Key Not Found
        return None

    flags = struct.unpack(">I", body[:extras_len])[0]
    value = body[extras_len + key_len:]

    return {
        'value': value,
        'flags': flags,
        'cas': cas,
        'opaque': opaque
    }

def _recv_exact(sock, n: int) -> bytes:
    data = b""
    while len(data) < n:
        chunk = sock.recv(n - len(data))
        if not chunk:
            raise ConnectionError("Connection closed")
        data += chunk
    return data

# GETQ — 静默获取
# 当 key 不存在时，不返回任何响应
# 适用于管道化请求：发送多个 GETQ，只有存在的 key 才返回响应

def binary_get_multi(sock, keys: list[str]) -> dict:
    """使用 GETKQ 批量获取"""
    for i, key in enumerate(keys):
        key_bytes = key.encode()
        header = struct.pack(
            "!BBHBBHIIQ",
            0x80, 0x0D, len(key_bytes), 0, 0x00, 0, len(key_bytes), i, 0
        )
        sock.sendall(header + key_bytes)

    # 发送 NOOP 作为结束标记
    noop_header = struct.pack(
        "!BBHBBHIIQ",
        0x80, 0x0A, 0, 0, 0x00, 0, 0, len(keys), 0
    )
    sock.sendall(noop_header)

    # 读取响应直到 NOOP 响应
    result = {}
    while True:
        resp_header = _recv_exact(sock, 24)
        (_, opcode, key_len, extras_len, _, status, body_len, opaque, cas) = \
            struct.unpack("!BBHBBHIIQ", resp_header)

        if opcode == 0x0A:  # NOOP
            break

        body = _recv_exact(sock, body_len)
        if status == 0x0000:
            flags = struct.unpack(">I", body[:extras_len])[0]
            key = body[extras_len:extras_len + key_len].decode()
            value = body[extras_len + key_len:]
            result[key] = {'value': value, 'flags': flags, 'cas': cas}

    return result

def binary_set(sock, key: str, value: bytes, flags: int = 0,
               exptime: int = 0, cas: int = 0,
               opcode: int = 0x01) -> dict:
    """
    二进制 SET/ADD/REPLACE 命令
    opcode: 0x01=SET, 0x02=ADD, 0x03=REPLACE
    """
    key_bytes = key.encode()
    extras = struct.pack(">II", flags, exptime)
    body_length = len(extras) + len(key_bytes) + len(value)

    header = struct.pack(
        "!BBHBBHIIQ",
        0x80,           # magic
        opcode,         # opcode
        len(key_bytes), # key length
        len(extras),    # extras length
        0x00,           # data type
        0,              # reserved
        body_length,    # body length
        0,              # opaque
        cas             # cas
    )

    sock.sendall(header + extras + key_bytes + value)

    # 读取响应
    resp_header = _recv_exact(sock, 24)
    (_, _, _, _, _, status, _, _, new_cas) = \
        struct.unpack("!BBHBBHIIQ", resp_header)

    return {
        'success': status == 0x0000,
        'status': status,
        'cas': new_cas
    }

def binary_delete(sock, key: str) -> bool:
    key_bytes = key.encode()
    header = struct.pack(
        "!BBHBBHIIQ",
        0x80, 0x04, len(key_bytes), 0, 0x00, 0, len(key_bytes), 0, 0
    )
    sock.sendall(header + key_bytes)

    resp_header = _recv_exact(sock, 24)
    (_, _, _, _, _, status, _, _, _) = \
        struct.unpack("!BBHBBHIIQ", resp_header)

    return status == 0x0000

def binary_incr_decr(sock, key: str, delta: int, initial: int = 0,
                     exptime: int = 0, incr: bool = True) -> int | None:
    key_bytes = key.encode()
    extras = struct.pack(">QQI", delta, initial, exptime)
    opcode = 0x05 if incr else 0x06

    header = struct.pack(
        "!BBHBBHIIQ",
        0x80, opcode, len(key_bytes), len(extras), 0x00, 0,
        len(extras) + len(key_bytes), 0, 0
    )

    sock.sendall(header + extras + key_bytes)

    resp_header = _recv_exact(sock, 24)
    (_, _, _, _, _, status, body_len, _, _) = \
        struct.unpack("!BBHBBHIIQ", resp_header)

    if status == 0x0001:  # Not Found
        return None

    body = _recv_exact(sock, body_len)
    result = struct.unpack(">Q", body)[0]
    return result

def binary_append(sock, key: str, value: bytes, prepend: bool = False) -> bool:
    key_bytes = key.encode()
    opcode = 0x0F if prepend else 0x0E

    header = struct.pack(
        "!BBHBBHIIQ",
        0x80, opcode, len(key_bytes), 0, 0x00, 0,
        len(key_bytes) + len(value), 0, 0
    )

    sock.sendall(header + key_bytes + value)

    resp_header = _recv_exact(sock, 24)
    (_, _, _, _, _, status, _, _, _) = \
        struct.unpack("!BBHBBHIIQ", resp_header)

    return status == 0x0000

def binary_flush(sock, delay: int = 0) -> bool:
    if delay > 0:
        extras = struct.pack(">I", delay)
        header = struct.pack(
            "!BBHBBHIIQ",
            0x80, 0x08, 0, 4, 0x00, 0, 4, 0, 0
        )
        sock.sendall(header + extras)
    else:
        header = struct.pack(
            "!BBHBBHIIQ",
            0x80, 0x08, 0, 0, 0x00, 0, 0, 0, 0
        )
        sock.sendall(header)

    resp_header = _recv_exact(sock, 24)
    (_, _, _, _, _, status, _, _, _) = \
        struct.unpack("!BBHBBHIIQ", resp_header)

    return status == 0x0000

def binary_stats(sock) -> dict[str, str]:
    header = struct.pack(
        "!BBHBBHIIQ",
        0x80, 0x10, 0, 0, 0x00, 0, 0, 0, 0
    )
    sock.sendall(header)

    result = {}
    while True:
        resp_header = _recv_exact(sock, 24)
        (_, _, key_len, _, _, status, body_len, _, _) = \
            struct.unpack("!BBHBBHIIQ", resp_header)

        if key_len == 0 and body_len == 0:
            break

        body = _recv_exact(sock, body_len)
        key = body[:key_len].decode()
        value = body[key_len:].decode()
        result[key] = value

    return result

def binary_noop(sock) -> bool:
    header = struct.pack(
        "!BBHBBHIIQ",
        0x80, 0x0A, 0, 0, 0x00, 0, 0, 0, 0
    )
    sock.sendall(header)

    resp_header = _recv_exact(sock, 24)
    (_, opcode, _, _, _, status, _, _, _) = \
        struct.unpack("!BBHBBHIIQ", resp_header)

    return opcode == 0x0A and status == 0x0000

def binary_multi_set(sock, items: list[tuple[str, bytes, int]]):
    """
    使用 SETQ 管道化批量写入
    items: [(key, value, exptime), ...]
    """
    for i, (key, value, exptime) in enumerate(items):
        key_bytes = key.encode()
        extras = struct.pack(">II", 0, exptime)
        body_len = len(extras) + len(key_bytes) + len(value)

        header = struct.pack(
            "!BBHBBHIIQ",
            0x80, 0x11, len(key_bytes), len(extras), 0x00, 0,
            body_len, i, 0
        )
        sock.sendall(header + extras + key_bytes + value)

    # 发送 NOOP 确认所有操作完成
    noop_header = struct.pack(
        "!BBHBBHIIQ",
        0x80, 0x0A, 0, 0, 0x00, 0, 0, len(items), 0
    )
    sock.sendall(noop_header)

    # 读取 NOOP 响应（所有 SETQ 成功则只有 NOOP 响应）
    while True:
        resp_header = _recv_exact(sock, 24)
        (_, opcode, _, _, _, _, _, opaque, _) = \
            struct.unpack("!BBHBBHIIQ", resp_header)

        if opcode == 0x0A:  # NOOP
            break

        # 如果收到非 NOOP 响应，说明某个 SETQ 失败
        body_len = struct.unpack("!I", resp_header[8:12])[0]
        if body_len > 0:
            _recv_exact(sock, body_len)

    return True

def binary_sasl_auth(sock, username: str, password: str) -> bool:
    """SASL PLAIN 认证"""
    # PLAIN 格式: \0username\0password
    challenge = b"\0" + username.encode() + b"\0" + password.encode()
    mechanism = b"PLAIN"

    header = struct.pack(
        "!BBHBBHIIQ",
        0x80, 0x21, len(mechanism), 0, 0x00, 0,
        len(mechanism) + len(challenge), 0, 0
    )
    sock.sendall(header + mechanism + challenge)

    resp_header = _recv_exact(sock, 24)
    (_, _, _, _, _, status, body_len, _, _) = \
        struct.unpack("!BBHBBHIIQ", resp_header)

    if body_len > 0:
        _recv_exact(sock, body_len)

    return status == 0x0000

#!/usr/bin/env python3
"""full_binary_client.py — 完整的 Memcached 二进制协议客户端"""

import struct
import socket
from typing import Optional

class FullBinaryMemcachedClient:
    def __init__(self, host='127.0.0.1', port=11211):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect((host, port))
        self._opaque = 0

    def _next_opaque(self) -> int:
        self._opaque += 1
        return self._opaque

    def _send_recv(self, opcode: int, key: bytes = b"",
                   value: bytes = b"", extras: bytes = b"",
                   cas: int = 0) -> dict:
        body_len = len(extras) + len(key) + len(value)
        header = struct.pack(
            "!BBHBBHIIQ",
            0x80, opcode, len(key), len(extras), 0x00, 0,
            body_len, self._next_opaque(), cas
        )
        self.sock.sendall(header + extras + key + value)

        resp = _recv_exact(self.sock, 24)
        (_, opcode, key_len, extras_len, _, status, body_len, opaque, cas) = \
            struct.unpack("!BBHBBHIIQ", resp)

        body = _recv_exact(self.sock, body_len) if body_len > 0 else b""
        return {
            'opcode': opcode, 'status': status,
            'key': body[extras_len:extras_len + key_len],
            'value': body[extras_len + key_len:],
            'extras': body[:extras_len], 'cas': cas, 'opaque': opaque
        }

    def version(self) -> str:
        r = self._send_recv(0x0B)
        return r['value'].decode()

    def set(self, key: str, value: bytes, flags: int = 0,
            exptime: int = 0, cas: int = 0) -> dict:
        extras = struct.pack(">II", flags, exptime)
        return self._send_recv(0x01, key.encode(), value, extras, cas)

    def add(self, key: str, value: bytes, flags: int = 0,
            exptime: int = 0) -> dict:
        extras = struct.pack(">II", flags, exptime)
        return self._send_recv(0x02, key.encode(), value, extras)

    def replace(self, key: str, value: bytes, flags: int = 0,
                exptime: int = 0) -> dict:
        extras = struct.pack(">II", flags, exptime)
        return self._send_recv(0x03, key.encode(), value, extras)

    def get(self, key: str) -> Optional[tuple[bytes, int, int]]:
        r = self._send_recv(0x00, key.encode())
        if r['status'] == 0x0001:
            return None
        flags = struct.unpack(">I", r['extras'])[0] if r['extras'] else 0
        return r['value'], flags, r['cas']

    def delete(self, key: str) -> bool:
        r = self._send_recv(0x04, key.encode())
        return r['status'] == 0x0000

    def incr(self, key: str, delta: int, initial: int = 0,
             exptime: int = 0) -> Optional[int]:
        extras = struct.pack(">QQI", delta, initial, exptime)
        r = self._send_recv(0x05, key.encode(), extras=extras)
        if r['status'] == 0x0001:
            return None
        return struct.unpack(">Q", r['value'])[0]

    def decr(self, key: str, delta: int, initial: int = 0,
             exptime: int = 0) -> Optional[int]:
        extras = struct.pack(">QQI", delta, initial, exptime)
        r = self._send_recv(0x06, key.encode(), extras=extras)
        if r['status'] == 0x0001:
            return None
        return struct.unpack(">Q", r['value'])[0]

    def flush(self, delay: int = 0) -> bool:
        if delay > 0:
            extras = struct.pack(">I", delay)
            r = self._send_recv(0x08, extras=extras)
        else:
            r = self._send_recv(0x08)
        return r['status'] == 0x0000

    def noop(self) -> bool:
        r = self._send_recv(0x0A)
        return r['status'] == 0x0000

    def close(self):
        self._send_recv(0x07)  # QUIT
        self.sock.close()


# 测试
client = FullBinaryMemcachedClient()
print(f"Version: {client.version()}")

client.set("bin:test", b"hello binary", flags=0, exptime=300)
result = client.get("bin:test")
if result:
    print(f"Value: {result[0].decode()}, Flags: {result[1]}, CAS: {result[2]}")

print(f"Incr: {client.incr('bin:counter', 1, initial=100)}")
print(f"Incr: {client.incr('bin:counter', 1)}")

client.delete("bin:test")
client.delete("bin:counter")
client.close()