import struct
import typing as t
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.kdf.concatkdf import ConcatKDFHash
from ..registry import Header
from ..util import to_bytes, urlsafe_b64decode


__all__ = [
    "derive_key_for_concat_kdf",
    "u32be_len_input",
]


def derive_key_for_concat_kdf(
        shared_key: bytes,
        header: Header,
        cek_size: int,
        key_size: t.Optional[int],
        tag: t.Optional[bytes] = None):
    # PartyUInfo
    apu_info = u32be_len_input(header.get("apu"), True)
    # PartyVInfo
    apv_info = u32be_len_input(header.get("apv"), True)
    # SuppPubInfo

    if key_size:
        alg_id = u32be_len_input(header["alg"])
        bit_size = key_size
    else:
        alg_id = u32be_len_input(header["enc"])
        bit_size = cek_size

    pub_info = struct.pack(">I", bit_size)
    fixed_info = alg_id + apu_info + apv_info + pub_info

    if tag:
        cctag = u32be_len_input(tag)
        fixed_info += cctag

    ckdf = ConcatKDFHash(
        algorithm=hashes.SHA256(),
        length=bit_size // 8,
        otherinfo=fixed_info,
        backend=default_backend(),
    )
    return ckdf.derive(shared_key)


def u32be_len_input(s, use_base64=False):
    if not s:
        return b"\x00\x00\x00\x00"
    if use_base64:
        s = urlsafe_b64decode(to_bytes(s))
    else:
        s = to_bytes(s)
    return struct.pack(">I", len(s)) + s
