from __future__ import annotations

import argparse
import struct
from pathlib import Path

from .xfalcon_blake_core import compress, digest_hex, words_hex
from .xfalcon_te import (
    XFALCON_PREFIX,
    xfalcon_te_hex_from_digest_hex,
    xfalcon_value_from_digest_hex,
)


MASK32 = 0xFFFFFFFF
INPUT_HEX_LEN = 80
PACKED_LEN = 64
M0_XOR_CONST = 0x3D
RAW_CHUNK_LEN = 256
FOLDED_BLOCK_LEN = 64
PARAM_BLOCK_BYTES = 0x01010020
CUSTOM_IV = [
    0xA92157F6, 0xA9A24FF4, 0x9138D3FD, 0x2A2193F3,
    0x2ADEF3F4, 0x9876EF16, 0x9ABED34F, 0x9103DE12,
]


def _validate_input_hex(input_hex: str) -> str:
    input_hex = input_hex.strip()
    if len(input_hex) != INPUT_HEX_LEN:
        raise ValueError(f"input_hex must be {INPUT_HEX_LEN} chars")
    if any(c not in "0123456789abcdefABCDEF" for c in input_hex):
        raise ValueError("input_hex must contain only hex chars")
    return input_hex


def _input_to_bytes(value: bytes | bytearray | str) -> bytes:
    if isinstance(value, str):
        return value.encode("utf-8")
    return bytes(value)


def _initial_h() -> list[int]:
    h = CUSTOM_IV.copy()
    h[0] ^= PARAM_BLOCK_BYTES
    return h


def _compress_v0(h: list[int], counter: int, final: bool) -> list[int]:
    v0 = h.copy() + CUSTOM_IV.copy()
    v0[12] ^= counter & MASK32
    if final:
        v0[14] ^= MASK32
    return v0


def _state_after_compress(h: list[int], block: bytes, counter: int, final: bool) -> list[int]:
    if len(block) != FOLDED_BLOCK_LEN:
        raise ValueError(f"folded block must be {FOLDED_BLOCK_LEN} bytes")
    m_words = [struct.unpack_from("<I", block, i * 4)[0] for i in range(16)]
    v_final = compress(_compress_v0(h, counter, final), m_words)
    return [(h[i] ^ v_final[i] ^ v_final[i + 8]) & MASK32 for i in range(8)]


def _chunk_logical_len(raw_len: int, offset: int) -> int:
    if offset >= raw_len:
        return 0
    return min(FOLDED_BLOCK_LEN, (raw_len - offset + 3) // 4)


def xfalcon_folded_block_from_raw_bytes(raw: bytes | bytearray, offset: int) -> bytes:
    """Fold one 256-byte VM raw chunk into the 64-byte BLAKE message block.

    The VM appends the fixed HUDR prefix to the caller input, then compresses
    each 256-byte raw chunk into 64 bytes by XORing four 64-byte lanes.
    """
    raw_bytes = bytes(raw)
    if offset < 0:
        raise ValueError("offset must be non-negative")

    block = bytearray(FOLDED_BLOCK_LEN)
    for j in range(FOLDED_BLOCK_LEN):
        value = 0
        for lane in range(4):
            pos = offset + lane * FOLDED_BLOCK_LEN + j
            if pos < len(raw_bytes):
                value ^= raw_bytes[pos]
        block[j] = value
    return bytes(block)


def xfalcon_folded_block_from_input_bytes(input_bytes: bytes | bytearray | str, offset: int) -> bytes:
    raw = _input_to_bytes(input_bytes) + XFALCON_PREFIX.encode("ascii")
    return xfalcon_folded_block_from_raw_bytes(raw, offset)


def xfalcon_message_words_from_folded_block(block: bytes | bytearray) -> list[int]:
    block_bytes = bytes(block)
    if len(block_bytes) != FOLDED_BLOCK_LEN:
        raise ValueError(f"folded block must be {FOLDED_BLOCK_LEN} bytes")
    return [struct.unpack_from("<I", block_bytes, i * 4)[0] for i in range(16)]


def xfalcon_digest_words_from_raw_bytes(raw: bytes | bytearray) -> list[int]:
    raw_bytes = bytes(raw)
    if not raw_bytes:
        raise ValueError("raw bytes must not be empty")

    h = _initial_h()
    counter = 0
    for offset in range(0, len(raw_bytes), RAW_CHUNK_LEN):
        block = xfalcon_folded_block_from_raw_bytes(raw_bytes, offset)
        counter += _chunk_logical_len(len(raw_bytes), offset)
        final = offset + RAW_CHUNK_LEN >= len(raw_bytes)
        h = _state_after_compress(h, block, counter, final)
    return h


def xfalcon_digest_words_from_input_bytes(input_bytes: bytes | bytearray | str) -> list[int]:
    raw = _input_to_bytes(input_bytes) + XFALCON_PREFIX.encode("ascii")
    return xfalcon_digest_words_from_raw_bytes(raw)


def xfalcon_digest_hex_from_input_bytes(input_bytes: bytes | bytearray | str) -> str:
    return digest_hex(xfalcon_digest_words_from_input_bytes(input_bytes))


def xfalcon_te_hex_from_input_bytes(input_bytes: bytes | bytearray | str) -> str:
    return xfalcon_te_hex_from_digest_hex(xfalcon_digest_hex_from_input_bytes(input_bytes))


def xfalcon_value_from_input_bytes(input_bytes: bytes | bytearray | str) -> str:
    return xfalcon_value_from_digest_hex(xfalcon_digest_hex_from_input_bytes(input_bytes))


def xfalcon_packed_block_from_input_hex(input_hex: str) -> bytes:
    """Rebuild the legacy pre-XOR packed block from the 80-char ASCII input."""
    input_hex = _validate_input_hex(input_hex)
    block = bytearray(xfalcon_folded_block_from_input_bytes(input_hex, 0))
    block[0] ^= M0_XOR_CONST
    return bytes(block)


def xfalcon_message_words_from_input_hex(input_hex: str) -> list[int]:
    """Rebuild BLAKE message words at x23+0x2fa8."""
    input_hex = _validate_input_hex(input_hex)
    return xfalcon_message_words_from_folded_block(
        xfalcon_folded_block_from_input_bytes(input_hex, 0)
    )


def xfalcon_digest_words_from_input_hex(input_hex: str) -> list[int]:
    input_hex = _validate_input_hex(input_hex)
    return xfalcon_digest_words_from_input_bytes(input_hex)


def xfalcon_digest_hex_from_input_hex(input_hex: str) -> str:
    return digest_hex(xfalcon_digest_words_from_input_hex(input_hex))


def xfalcon_te_hex_from_input_hex(input_hex: str) -> str:
    return xfalcon_te_hex_from_digest_hex(xfalcon_digest_hex_from_input_hex(input_hex))


def xfalcon_value_from_input_hex(input_hex: str) -> str:
    return xfalcon_value_from_digest_hex(xfalcon_digest_hex_from_input_hex(input_hex))


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("input_hex", nargs="?", help="80-char ASCII hex input passed to __NS_xfalcon")
    source_group = parser.add_mutually_exclusive_group()
    source_group.add_argument("--file", help="read arbitrary xfalcon input bytes from file")
    source_group.add_argument("--text", help="use arbitrary UTF-8 xfalcon input text")
    group = parser.add_mutually_exclusive_group()
    group.add_argument("--digest", action="store_true", help="print digest hex only")
    group.add_argument("--te", action="store_true", help="print $TE_ raw hex only")
    group.add_argument("--value", action="store_true", help="print full __NS_xfalcon value only")
    group.add_argument("--packed", action="store_true", help="print direct legacy packed block hex only")
    group.add_argument("--m", action="store_true", help="print first folded BLAKE message words only")
    args = parser.parse_args()

    if args.file:
        input_bytes = Path(args.file).read_bytes()
        digest = xfalcon_digest_hex_from_input_bytes(input_bytes)
        first_block = xfalcon_folded_block_from_input_bytes(input_bytes, 0)
    elif args.text is not None:
        digest = xfalcon_digest_hex_from_input_bytes(args.text)
        first_block = xfalcon_folded_block_from_input_bytes(args.text, 0)
    else:
        if not args.input_hex:
            parser.error("input_hex is required unless --file or --text is used")
        digest = xfalcon_digest_hex_from_input_hex(args.input_hex)
        first_block = xfalcon_folded_block_from_input_bytes(args.input_hex, 0)

    if args.digest:
        print(digest)
    elif args.te:
        print(xfalcon_te_hex_from_digest_hex(digest))
    elif args.value:
        print(xfalcon_value_from_digest_hex(digest))
    elif args.packed:
        if args.file or args.text is not None:
            parser.error("--packed is only defined for direct 80-char ASCII hex input")
        print(xfalcon_packed_block_from_input_hex(args.input_hex).hex())
    elif args.m:
        print(words_hex(xfalcon_message_words_from_folded_block(first_block)))
    else:
        print(f"digest={digest}")
        print(f"te={xfalcon_te_hex_from_digest_hex(digest)}")
        print(f"value={xfalcon_value_from_digest_hex(digest)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())