# import pyarrow.flight as flight
# import base64
#
#
# class SimpleAuthHandler(flight.ServerAuthHandler):
#     def __init__(self):
#         # 假设这是一个硬编码的用户名/密码
#         self._valid_credentials = {"admin": "admin"}
#
#     def authenticate(self,  outgoing: flight.ServerAuthSender):
#         # 接收客户端发来的 credentials 对象（序列化后的字节数组）
#         # credentials_bytes = incoming.read()
#
#         # 反序列化（你需要和客户端用相同的序列化协议）
#         credentials_str = credentials_bytes.decode('utf-8')
#         username, password = credentials_str.split(":")  # 假设简单格式为 "user:pass"
#
#         if self._valid_credentials.get(username) == password:
#             # 认证成功，发送 session token
#             token = f"token-for-{username}"
#             outgoing.write(token.encode('utf-8'))
#             print(f"Server: Authenticated {username}")
#         else:
#             raise flight.FlightUnauthenticatedError("Invalid credentials")
#
#     def is_valid(self, token: bytes):
#         # 校验 token 是否有效。这里只是示例，真实场景可能查数据库或缓存
#         token_str = token.decode('utf-8')
#         if token_str.startswith("token-for-"):
#             username = token_str[len("token-for-"):]
#             return username
#         raise flight.FlightUnauthenticatedError("Invalid token")
import struct
import pyarrow.flight as flight

class PairAuthHandler(flight.ClientAuthHandler):
    def __init__(self, username: str, password: str):
        self.username = username
        self.password = password
        self._token = None

    def authenticate(self, outgoing, incoming):
        user_bytes = self.username.encode("utf-8")
        pwd_bytes = self.password.encode("utf-8")

        # 按照 Scala 的 encodePair 格式拼接

        payload = (
            struct.pack('<B', 1) +
            struct.pack(">I", len(user_bytes)) + user_bytes +
            struct.pack(">I", len(pwd_bytes)) + pwd_bytes
        )

        outgoing.write(payload)

        # 读取服务端返回的 token
        self._token = incoming.read()

    def get_token(self):
        return self._token

