"""H5 `__NS_sig3` 34-byte envelope helpers.

The Nebula H5 bridge returns a 68-hex digest shape that is distinct from the
regular API 10418 48-hex `__NS_sig3`. The final byte is a checksum mask; the
first 33 bytes are XOR-mixed with `(mask ^ index)`.
"""

from __future__ import annotations

import zlib
from dataclasses import dataclass

from .sig3 import KWSG_10418_HMAC_KEY, kwsg_10418_binary48, load_kwsg_10418_tables


H5_SIG3_LENGTH = 34
H5_SIG3_HEX_LENGTH = H5_SIG3_LENGTH * 2
H5_SIG3_MAGIC = bytes.fromhex("54450130")
H5_SIG3_BLOCK_TAG = bytes.fromhex("9f01")
H5_SIG3_DEFAULT_SESSION_SEED = 0x4B4B4BD1
H5_SIG3_DEFAULT_TAIL = bytes.fromhex("0100000000")
H5_ENCODE_SHA_IV1 = (
    0xE066DF2C,
    0xA9888E01,
    0xB18EE43D,
    0x501AFA25,
    0x53A550E5,
    0xA4BCE311,
    0x70E554F5,
    0x6AA046EE,
)
H5_ENCODE_SHA_IV2 = (
    0x3C66483A,
    0x0096038E,
    0xD62FEC37,
    0x346D7B3D,
    0xDE7FB319,
    0xF56D5DD5,
    0xE614E937,
    0x6DD5E338,
)
H5_ENCODE_SHA_K = (
    0x428A2F98,
    0x71374491,
    0xB5C0FBCF,
    0xE9B5DBA5,
    0x3956C25B,
    0x59F111F1,
    0x923F82A4,
    0xAB1C5ED5,
    0xD807AA98,
    0x12835B01,
    0x243185BE,
    0x550C7DC3,
    0x72BE5D74,
    0x80DEB1FE,
    0x9BDC06A7,
    0xC19BF174,
    0xE49B69C1,
    0xEFBE4786,
    0x0FC19DC6,
    0x240CA1CC,
    0x2DE92C6F,
    0x4A7484AA,
    0x5CB0A9DC,
    0x76F988DA,
    0x983E5152,
    0xA831C66D,
    0xB00327C8,
    0xBF597FC7,
    0xC6E00BF3,
    0xD5A79147,
    0x06CA6351,
    0x14292967,
    0x27B70A85,
    0x2E1B2138,
    0x4D2C6DFC,
    0x53380D13,
    0x650A7354,
    0x766A0ABB,
    0x81C2C92E,
    0x92722C85,
    0xA2BFE8A1,
    0xA81A664B,
    0xC24B8B70,
    0xC76C51A3,
    0xD192E819,
    0xD6990624,
    0xF40E3585,
    0x106AA070,
    0x19A4C116,
    0x1E376C08,
    0x2748774C,
    0x34B0BCB5,
    0x391C0CB3,
    0x4ED8AA4A,
    0x5B9CCA4F,
    0x682E6FF3,
    0x748F82EE,
    0x78A5636F,
    0x84C87814,
    0x8CC70208,
    0x90BEFFFA,
    0xA4506CEB,
    0xBEF9A3F7,
    0xC67178F2,
)


@dataclass(frozen=True)
class H5Sig3Fields:
    """扰动前的 H5 sig3 字段视图。"""

    raw_hex: str
    preimage: bytes
    mask: int
    mask_ok: bool
    magic_ok: bool
    first_tag_ok: bool
    second_tag_ok: bool
    tail_ok: bool
    session_seed: int
    counter: int
    crc32: int
    elapsed_ms: int
    state_value: int
    tail: bytes

    @property
    def ok(self) -> bool:
        return (
            self.mask_ok
            and self.magic_ok
            and self.first_tag_ok
            and self.second_tag_ok
            and self.tail_ok
        )


def _coerce_digest(value: str | bytes | bytearray) -> bytes:
    if isinstance(value, str):
        if len(value) != H5_SIG3_HEX_LENGTH:
            raise ValueError("H5 sig3 must be exactly 68 hex chars")
        return bytes.fromhex(value)
    digest = bytes(value)
    if len(digest) != H5_SIG3_LENGTH:
        raise ValueError("H5 sig3 must be exactly 34 bytes")
    return digest


def h5_sig3_expected_mask(preimage33: bytes | bytearray) -> int:
    """按已观测到的 H5 envelope 规则计算最后 1 字节 mask。"""
    if len(preimage33) != H5_SIG3_LENGTH - 1:
        raise ValueError("H5 sig3 preimage prefix must be exactly 33 bytes")
    return (-sum(preimage33)) & 0xFF


def h5_sig3_unmix(value: str | bytes | bytearray) -> tuple[bytes, int, bool]:
    """恢复扰动前 34 字节，其中最后 1 字节固定置 0。"""
    digest = _coerce_digest(value)
    mask = digest[-1]
    preimage = bytearray(H5_SIG3_LENGTH)
    for i in range(H5_SIG3_LENGTH - 1):
        preimage[i] = digest[i] ^ ((mask ^ i) & 0xFF)
    expected = h5_sig3_expected_mask(preimage[: H5_SIG3_LENGTH - 1])
    return bytes(preimage), mask, mask == expected


def h5_sig3_mix(preimage: bytes | bytearray, mask: int | None = None) -> bytes:
    """由扰动前字段重建 34-byte H5 sig3 digest。"""
    pre = bytearray(preimage)
    if len(pre) == H5_SIG3_LENGTH - 1:
        pre.append(0)
    if len(pre) != H5_SIG3_LENGTH:
        raise ValueError("H5 sig3 preimage must be exactly 33 or 34 bytes")
    if mask is None:
        mask = h5_sig3_expected_mask(pre[: H5_SIG3_LENGTH - 1])
    mask &= 0xFF
    out = bytearray(H5_SIG3_LENGTH)
    for i in range(H5_SIG3_LENGTH - 1):
        out[i] = pre[i] ^ ((mask ^ i) & 0xFF)
    out[-1] = mask
    return bytes(out)


def parse_h5_sig3(value: str | bytes | bytearray) -> H5Sig3Fields:
    """解析 H5 68hex `__NS_sig3` 的 envelope 字段。"""
    digest = _coerce_digest(value)
    pre, mask, mask_ok = h5_sig3_unmix(digest)
    return H5Sig3Fields(
        raw_hex=digest.hex(),
        preimage=pre,
        mask=mask,
        mask_ok=mask_ok,
        magic_ok=pre[:4] == H5_SIG3_MAGIC,
        first_tag_ok=pre[8:10] == H5_SIG3_BLOCK_TAG,
        second_tag_ok=pre[22:24] == H5_SIG3_BLOCK_TAG,
        tail_ok=pre[28:33] == H5_SIG3_DEFAULT_TAIL,
        session_seed=int.from_bytes(pre[4:8], "little"),
        counter=int.from_bytes(pre[10:14], "little"),
        crc32=int.from_bytes(pre[14:18], "little"),
        elapsed_ms=int.from_bytes(pre[18:22], "little"),
        state_value=int.from_bytes(pre[24:28], "little"),
        tail=pre[28:33],
    )


def build_h5_sig3_preimage(
    *,
    crc32_value: int,
    counter: int,
    elapsed_ms: int,
    state_value: int,
    session_seed: int = H5_SIG3_DEFAULT_SESSION_SEED,
    tail: bytes = H5_SIG3_DEFAULT_TAIL,
) -> bytes:
    """构造 H5 sig3 扰动前字段，便于回放样本和后续接入真实 CRC。"""
    if len(tail) != 5:
        raise ValueError("H5 sig3 tail must be exactly 5 bytes")
    pre = bytearray(H5_SIG3_LENGTH)
    pre[:4] = H5_SIG3_MAGIC
    pre[4:8] = (int(session_seed) & 0xFFFFFFFF).to_bytes(4, "little")
    pre[8:10] = H5_SIG3_BLOCK_TAG
    pre[10:14] = (int(counter) & 0xFFFFFFFF).to_bytes(4, "little")
    pre[14:18] = (int(crc32_value) & 0xFFFFFFFF).to_bytes(4, "little")
    pre[18:22] = (int(elapsed_ms) & 0xFFFFFFFF).to_bytes(4, "little")
    pre[22:24] = H5_SIG3_BLOCK_TAG
    pre[24:28] = (int(state_value) & 0xFFFFFFFF).to_bytes(4, "little")
    pre[28:33] = tail
    return bytes(pre)


def h5_sig3_from_fields(
    *,
    crc32_value: int,
    counter: int,
    elapsed_ms: int,
    state_value: int,
    session_seed: int = H5_SIG3_DEFAULT_SESSION_SEED,
    tail: bytes = H5_SIG3_DEFAULT_TAIL,
) -> str:
    """按字段直接生成 H5 68hex sig3。"""
    return h5_sig3_mix(
        build_h5_sig3_preimage(
            crc32_value=crc32_value,
            counter=counter,
            elapsed_ms=elapsed_ms,
            state_value=state_value,
            session_seed=session_seed,
            tail=tail,
        )
    ).hex()


def _rotr32(value: int, bits: int) -> int:
    value &= 0xFFFFFFFF
    return ((value >> bits) | (value << (32 - bits))) & 0xFFFFFFFF


def _sha256_compress_block(state: tuple[int, ...], block64: bytes) -> tuple[int, ...]:
    if len(block64) != 64:
        raise ValueError("SHA-256 compression block must be exactly 64 bytes")
    w = [int.from_bytes(block64[i : i + 4], "big") for i in range(0, 64, 4)]
    for i in range(16, 64):
        s0 = _rotr32(w[i - 15], 7) ^ _rotr32(w[i - 15], 18) ^ (w[i - 15] >> 3)
        s1 = _rotr32(w[i - 2], 17) ^ _rotr32(w[i - 2], 19) ^ (w[i - 2] >> 10)
        w.append((w[i - 16] + s0 + w[i - 7] + s1) & 0xFFFFFFFF)

    a, b, c, d, e, f, g, h = (item & 0xFFFFFFFF for item in state)
    for i in range(64):
        s1 = _rotr32(e, 6) ^ _rotr32(e, 11) ^ _rotr32(e, 25)
        ch = (e & f) ^ ((~e) & g)
        temp1 = (h + s1 + ch + H5_ENCODE_SHA_K[i] + w[i]) & 0xFFFFFFFF
        s0 = _rotr32(a, 2) ^ _rotr32(a, 13) ^ _rotr32(a, 22)
        maj = (a & b) ^ (a & c) ^ (b & c)
        temp2 = (s0 + maj) & 0xFFFFFFFF
        h = g
        g = f
        f = e
        e = (d + temp1) & 0xFFFFFFFF
        d = c
        c = b
        b = a
        a = (temp1 + temp2) & 0xFFFFFFFF

    return (
        (state[0] + a) & 0xFFFFFFFF,
        (state[1] + b) & 0xFFFFFFFF,
        (state[2] + c) & 0xFFFFFFFF,
        (state[3] + d) & 0xFFFFFFFF,
        (state[4] + e) & 0xFFFFFFFF,
        (state[5] + f) & 0xFFFFFFFF,
        (state[6] + g) & 0xFFFFFFFF,
        (state[7] + h) & 0xFFFFFFFF,
    )


def _sha256_continue_from_virtual_prefix(
    data: bytes,
    state: tuple[int, ...],
    prefix_len: int = 64,
) -> bytes:
    """从已压缩过 `prefix_len` 字节的自定义 SHA-256 state 继续计算。"""
    if len(state) != 8:
        raise ValueError("SHA-256 state must contain exactly 8 words")
    total_len = prefix_len + len(data)
    padded = bytearray(data)
    padded.append(0x80)
    while (prefix_len + len(padded)) % 64 != 56:
        padded.append(0)
    padded += (total_len * 8).to_bytes(8, "big")

    current = tuple(word & 0xFFFFFFFF for word in state)
    for offset in range(0, len(padded), 64):
        current = _sha256_compress_block(current, bytes(padded[offset : offset + 64]))
    return b"".join(word.to_bytes(4, "big") for word in current)


def h5_encode_sha_digest(sign_input: str | bytes | bytearray) -> bytes:
    """复现 H5 `$encode` 中用于 envelope CRC 字段的双层 SHA digest。"""
    if isinstance(sign_input, str):
        data = sign_input.encode("utf-8")
    else:
        data = bytes(sign_input)
    first = _sha256_continue_from_virtual_prefix(data, H5_ENCODE_SHA_IV1, 64)
    return _sha256_continue_from_virtual_prefix(first, H5_ENCODE_SHA_IV2, 64)


def h5_sig3_crc32_from_sign_input(sign_input: str | bytes | bytearray) -> int:
    """返回 H5 envelope 14..17 字节对应的 little-endian 字段值。

    名字保留 `crc32` 是为了兼容既有字段命名；实际已确认不是标准 CRC32，
    而是 `$encode` 自定义双层 SHA digest 的首 32-bit word（按 envelope
    小端字段视图解释）。
    """
    digest = h5_encode_sha_digest(sign_input)
    return int.from_bytes(digest[:4], "little")


def h5_sig3_crc32_from_10418_input(
    input_value: str | bytes | bytearray,
    t1: bytes | None = None,
    t2: bytes | None = None,
    hmac_key: bytes = KWSG_10418_HMAC_KEY,
) -> int:
    """候选 CRC 生成器：CRC32(10418 binary48(input))。

    这个 helper 用于离线枚举 H5 `secPlain` 形态；当前只表示一个已知
    强候选路径，不代表 H5 bridge 的输入已经完全闭合。
    """
    if isinstance(input_value, str):
        input_bytes = input_value.encode("utf-8")
    else:
        input_bytes = bytes(input_value)
    if t1 is None and t2 is None:
        t1, t2 = load_kwsg_10418_tables()
    elif t1 is None or t2 is None:
        raise ValueError("t1 and t2 must be provided together")
    return zlib.crc32(kwsg_10418_binary48(input_bytes, t1, t2, hmac_key)) & 0xFFFFFFFF


__all__ = [
    "H5_SIG3_BLOCK_TAG",
    "H5_SIG3_DEFAULT_SESSION_SEED",
    "H5_SIG3_DEFAULT_TAIL",
    "H5_SIG3_HEX_LENGTH",
    "H5_SIG3_LENGTH",
    "H5_SIG3_MAGIC",
    "H5Sig3Fields",
    "build_h5_sig3_preimage",
    "h5_sig3_crc32_from_10418_input",
    "h5_sig3_crc32_from_sign_input",
    "h5_sig3_expected_mask",
    "h5_encode_sha_digest",
    "h5_sig3_from_fields",
    "h5_sig3_mix",
    "h5_sig3_unmix",
    "parse_h5_sig3",
]