"""
Copyright 2025 Beijing Volcano Engine Technology Co., Ltd. All Rights Reserved.
SPDX-license-identifier: BSD-3-Clause

Migrated from Server/token.js
"""

import base64
import hashlib
import hmac
import random
import struct
import time

VERSION = "001"
VERSION_LENGTH = 3
APP_ID_LENGTH = 24

privileges = {
    "PrivPublishStream": 0,
    "privPublishAudioStream": 1,
    "privPublishVideoStream": 2,
    "privPublishDataStream": 3,
    "PrivSubscribeStream": 4,
}


class ByteBuf:
    def __init__(self):
        self._buf = bytearray()

    def put_uint16(self, v: int) -> "ByteBuf":
        self._buf += struct.pack("<H", v)
        return self

    def put_uint32(self, v: int) -> "ByteBuf":
        self._buf += struct.pack("<I", v)
        return self

    def put_bytes(self, b: bytes) -> "ByteBuf":
        self.put_uint16(len(b))
        self._buf += b
        return self

    def put_string(self, s: str) -> "ByteBuf":
        return self.put_bytes(s.encode("utf-8"))

    def put_tree_map_uint32(self, m: dict) -> "ByteBuf":
        if not m:
            self.put_uint16(0)
            return self
        sorted_items = sorted(m.items(), key=lambda x: int(x[0]))
        self.put_uint16(len(sorted_items))
        for key, value in sorted_items:
            self.put_uint16(int(key))
            self.put_uint32(int(value))
        return self

    def pack(self) -> bytes:
        return bytes(self._buf)


def _encode_hmac(key: str, message: bytes) -> bytes:
    return hmac.new(key.encode("utf-8"), message, hashlib.sha256).digest()


class AccessToken:
    def __init__(self, app_id: str, app_key: str, room_id: str, user_id: str):
        self.app_id = app_id
        self.app_key = app_key
        self.room_id = room_id
        self.user_id = user_id
        self.issued_at = int(time.time())
        random.seed(time.time())
        self.nonce = random.randint(1, 99999999)
        self.expire_at = 0
        self._privileges: dict = {}

    def add_privilege(self, privilege: int, expire_timestamp: int):
        self._privileges[privilege] = expire_timestamp
        if privilege == privileges["PrivPublishStream"]:
            self._privileges[privileges["privPublishVideoStream"]] = expire_timestamp
            self._privileges[privileges["privPublishAudioStream"]] = expire_timestamp
            self._privileges[privileges["privPublishDataStream"]] = expire_timestamp

    def expire_time(self, expire_timestamp: int):
        self.expire_at = expire_timestamp

    def _pack_msg(self) -> bytes:
        buf = ByteBuf()
        buf.put_uint32(self.nonce)
        buf.put_uint32(self.issued_at)
        buf.put_uint32(self.expire_at)
        buf.put_string(self.room_id)
        buf.put_string(self.user_id)
        buf.put_tree_map_uint32(self._privileges)
        return buf.pack()

    def serialize(self) -> str:
        msg = self._pack_msg()
        signature = _encode_hmac(self.app_key, msg)
        content = ByteBuf().put_bytes(msg).put_bytes(signature).pack()
        return VERSION + self.app_id + base64.b64encode(content).decode("utf-8")