from Crypto.Cipher import AES
import hashlib
import struct

# Example implementation updated for the full signed int32 legacy fid domain.
# This keeps the algorithm reversible for all observed NavSea legacy fid values.
MIN_VAL = -(1 << 31)
MAX_VAL = (1 << 31) - 1
N = MAX_VAL - MIN_VAL + 1  # 4294967296

ROUNDS = 10
DOMAIN_BITS = 32
DOMAIN_SIZE = 1 << DOMAIN_BITS


def _prf(key: bytes, round_no: int, value: int) -> int:
    """
    AES-based PRF, returns 16-bit integer.
    """
    cipher = AES.new(key, AES.MODE_ECB)
    data = struct.pack(">IQ", round_no, value)
    block = hashlib.sha256(data).digest()[:16]
    out = cipher.encrypt(block)
    return int.from_bytes(out[-2:], "big")


def _feistel_permute(x: int, key: bytes, encrypt: bool) -> int:
    """
    32-bit Feistel permutation.
    """
    if not (0 <= x < DOMAIN_SIZE):
        raise ValueError("x out of 32-bit domain")

    l = (x >> 16) & 0xFFFF
    r = x & 0xFFFF

    rounds = range(ROUNDS) if encrypt else reversed(range(ROUNDS))

    for rnd in rounds:
        if encrypt:
            f = _prf(key, rnd, r)
            l, r = r, l ^ f
        else:
            f = _prf(key, rnd, l)
            l, r = r ^ f, l

    return ((l << 16) | r) & 0xFFFFFFFF


def encrypt_number(x: int, key: bytes) -> int:
    """
    Encrypt x in [-2147483648, 2147483647]
    Return non-negative integer in [0, N-1]
    """
    if len(key) not in (16, 24, 32):
        raise ValueError("AES key must be 16, 24, or 32 bytes")
    if not (MIN_VAL <= x <= MAX_VAL):
        raise ValueError("input out of range")

    # Map original range to [0, N-1]
    r = x - MIN_VAL

    # Cycle-walking into [0, N-1]
    while True:
        r = _feistel_permute(r, key, True)
        if r < N:
            return r


def decrypt_number(y: int, key: bytes) -> int:
    """
    Decrypt non-negative integer y in [0, N-1]
    Return original integer in [-2147483648, 2147483647]
    """
    if len(key) not in (16, 24, 32):
        raise ValueError("AES key must be 16, 24, or 32 bytes")
    if not (0 <= y < N):
        raise ValueError("encrypted value out of range")

    r = y

    # Reverse cycle-walking
    while True:
        r = _feistel_permute(r, key, False)
        if r < N:
            return r + MIN_VAL


if __name__ == "__main__":
    key = b"example-secret!!"  # 16 bytes

    samples = [
        -2147483648,
        -1,
        0,
        1,
        12345678,
        2147483647,
    ]

    for x in samples:
        enc = encrypt_number(x, key)
        dec = decrypt_number(enc, key)
        print(f"x={x}, enc={enc}, dec={dec}, ok={x == dec}")
