"""HKDF — HKDF-SHA256 key derivation.

Ported from net.i2p.crypto.HKDF.
Implements RFC 5869 with HMAC-SHA256.
One or two 32-byte outputs, with or without info.
"""

import hmac as _hmac
import hashlib


class HKDF:
    """HKDF-SHA256 key derivation. Thread-safe, no state."""

    def __init__(self) -> None:
        pass

    def calculate(self, key: bytes, data: bytes,
                  info: str = "",
                  out: bytearray | None = None,
                  out2: bytearray | None = None,
                  off2: int = 0) -> bytes:
        """HKDF-SHA256 extract-then-expand.

        Args:
            key: first 32 bytes used as the salt for extract
            data: input keying material
            info: optional context string (ASCII)
            out: if provided, 32-byte output buffer (output 1)
            out2: if provided, 32-byte output buffer (output 2)
            off2: offset into out2

        Returns:
            First 32-byte output (T1). If out2 is provided, also writes T2.
        """
        k = key[:32]

        # Extract: PRK = HMAC-SHA256(salt=key, data)
        prk = _hmac.new(k, data, hashlib.sha256).digest()

        # Expand T1: HMAC-SHA256(PRK, info || 0x01)
        info_bytes = info.encode("ascii") if info else b""
        t1 = _hmac.new(prk, info_bytes + b"\x01", hashlib.sha256).digest()

        if out is not None:
            out[0:32] = t1

        if out2 is not None:
            # Expand T2: HMAC-SHA256(PRK, T1 || info || 0x02)
            t2 = _hmac.new(prk, t1 + info_bytes + b"\x02", hashlib.sha256).digest()
            out2[off2 : off2 + 32] = t2

        return t1

    def extract_and_expand(self, salt: bytes, ikm: bytes, info: bytes, length: int) -> bytes:
        """General-purpose HKDF-SHA256 per RFC 5869.

        Args:
            salt: salt bytes
            ikm: input keying material
            info: context/application-specific info
            length: output length in bytes (max 255 * 32 = 8160)

        Returns:
            Derived key material of requested length.
        """
        # Extract
        prk = _hmac.new(salt, ikm, hashlib.sha256).digest()

        # Expand
        n = (length + 31) // 32
        okm = b""
        t = b""
        for i in range(1, n + 1):
            t = _hmac.new(prk, t + info + bytes([i]), hashlib.sha256).digest()
            okm += t

        return okm[:length]