"""Optional Numba fast path for `akira_decrypt.py recover`, stage 1 (the KCipher-2 seed-pair test).

akira_decrypt.py stays standard-library only; it imports this module only when numpy and numba are
installed (pip install numpy numba) and falls back to its pure-Python matcher otherwise.

The kernel is a line-for-line port of akira_decrypt._k2_ik / _k2_match. All arithmetic is done in
int64 and masked to 32 bits exactly as the Python reference does, so results are bit-identical;
selftest() checks that against the reference on random keys/IVs before anything relies on it.
"""
import numpy as np
from numba import njit, prange, get_num_threads, set_num_threads

import akira_decrypt as ad

M32 = 0xFFFFFFFF
_T = np.array([ad._T0, ad._T1, ad._T2, ad._T3], dtype=np.int64)
_A = np.array([ad.AMUL0, ad.AMUL1, ad.AMUL2, ad.AMUL3], dtype=np.int64)


@njit(cache=True, inline="always")
def _sub(T, x):
    return T[0, x & 255] ^ T[1, (x >> 8) & 255] ^ T[2, (x >> 16) & 255] ^ T[3, x >> 24]


@njit(cache=True, inline="always")
def _match(T, A, ik, v0, v1, v2, v3, target, tmask):
    """KCipher-2 init with key schedule ik and IV (v0..v3); True if the keystream words match target."""
    a0, a1, a2, a3, a4 = ik[4], ik[3], ik[2], ik[1], ik[0]
    b0, b1, b2, b3, b4, b5, b6, b7, b8, b9, b10 = ik[10], ik[11], v0, v1, ik[8], ik[9], v2, v3, ik[7], ik[5], ik[6]
    l1 = r1 = l2 = r2 = np.int64(0)
    for _ in range(24):
        nl1 = _sub(T, (r2 + b4) & M32)
        nr1 = _sub(T, (l2 + b9) & M32)
        nl2 = _sub(T, l1)
        nr2 = _sub(T, r1)
        na4 = ((a0 << 8) & M32) ^ A[0, a0 >> 24] ^ a3 ^ ((b0 + r2) & M32) ^ r1 ^ a4
        if (a2 >> 30) & 1:
            t1 = ((b0 << 8) & M32) ^ A[1, b0 >> 24]
        else:
            t1 = ((b0 << 8) & M32) ^ A[2, b0 >> 24]
        if a2 >> 31:
            t2 = ((b8 << 8) & M32) ^ A[3, b8 >> 24]
        else:
            t2 = b8
        nb10 = t1 ^ b1 ^ b6 ^ t2 ^ ((b10 + l2) & M32) ^ l1 ^ a0
        a0, a1, a2, a3, a4 = a1, a2, a3, a4, na4
        b0, b1, b2, b3, b4, b5, b6, b7, b8, b9, b10 = b1, b2, b3, b4, b5, b6, b7, b8, b9, b10, nb10
        l1, r1, l2, r2 = nl1, nr1, nl2, nr2
    n = target.shape[0]
    i = 0
    while True:
        if ((((b10 + l2) & M32) ^ l1 ^ a0) & tmask[i]) != target[i]:
            return False
        i += 1
        if i == n:
            return True
        if ((((b0 + r2) & M32) ^ r1 ^ a4) & tmask[i]) != target[i]:
            return False
        i += 1
        if i == n:
            return True
        nl1 = _sub(T, (r2 + b4) & M32)
        nr1 = _sub(T, (l2 + b9) & M32)
        nl2 = _sub(T, l1)
        nr2 = _sub(T, r1)
        na4 = ((a0 << 8) & M32) ^ A[0, a0 >> 24] ^ a3
        if (a2 >> 30) & 1:
            t1 = ((b0 << 8) & M32) ^ A[1, b0 >> 24]
        else:
            t1 = ((b0 << 8) & M32) ^ A[2, b0 >> 24]
        if a2 >> 31:
            t2 = ((b8 << 8) & M32) ^ A[3, b8 >> 24]
        else:
            t2 = b8
        nb10 = t1 ^ b1 ^ b6 ^ t2
        a0, a1, a2, a3, a4 = a1, a2, a3, a4, na4
        b0, b1, b2, b3, b4, b5, b6, b7, b8, b9, b10 = b1, b2, b3, b4, b5, b6, b7, b8, b9, b10, nb10
        l1, r1, l2, r2 = nl1, nr1, nl2, nr2


@njit(cache=True, inline="always")
def _key_schedule(T, k0, k1, k2, k3, ik):
    ik[0], ik[1], ik[2], ik[3] = k0, k1, k2, k3
    ik[4] = k0 ^ _sub(T, ((k3 << 8) | (k3 >> 24)) & M32) ^ 0x01000000
    ik[5] = ik[1] ^ ik[4]; ik[6] = ik[2] ^ ik[5]; ik[7] = ik[3] ^ ik[6]
    ik[8] = ik[4] ^ _sub(T, ((ik[7] << 8) | (ik[7] >> 24)) & M32) ^ 0x02000000
    ik[9] = ik[5] ^ ik[8]; ik[10] = ik[6] ^ ik[9]; ik[11] = ik[7] ^ ik[10]


@njit(parallel=True, cache=True)
def _scan(words, kidx, j0, cnt, d0, T, A, target, tmask):
    """For key candidate i (row kidx[i] of words), test IV rows j0[i]+d for d in d0[i]..cnt[i]-1.
    Returns, per key candidate, the first matching d or -1."""
    n = kidx.shape[0]
    out = np.full(n, -1, np.int64)
    for i in prange(n):
        ik = np.empty(12, np.int64)
        k = kidx[i]
        _key_schedule(T, words[k, 0], words[k, 1], words[k, 2], words[k, 3], ik)
        base = j0[i]
        for d in range(d0[i], cnt[i]):
            j = base + d
            if _match(T, A, ik, words[j, 0], words[j, 1], words[j, 2], words[j, 3], target, tmask):
                out[i] = d
                break
    return out


def words_from_table(buf, count):
    """Yarrow outputs (count x 16 bytes, as stored in the shared table) -> (count, 4) int64 big-endian words."""
    return np.frombuffer(buf, dtype=">u4", count=4 * count).reshape(count, 4).astype(np.int64)


def scan(words, kidx, j0, cnt, target, tmask, d0=None):
    if d0 is None:
        d0 = np.zeros_like(cnt)
    return _scan(words, kidx, j0, cnt, d0, _T, _A, np.asarray(target, np.int64), np.asarray(tmask, np.int64))


def threads(n=None):
    if n:
        set_num_threads(min(n, get_num_threads()))
    return get_num_threads()


def selftest(trials=2000, seed=1):
    """Fast kernel vs the pure-Python reference on random keys/IVs (match + mismatch, full and masked words)."""
    import random, struct
    rnd = random.Random(seed)
    bad = 0
    for t in range(trials):
        key, iv = rnd.randbytes(16), rnd.randbytes(16)
        ks = ad.KCipher2(key, iv).keystream(16)
        nbytes = (16, 7, 8, 13)[t % 4]
        nw = (nbytes + 3) // 4
        tm = [M32] * nw
        if nbytes % 4:
            tm[-1] = (M32 << (8 * (4 - nbytes % 4))) & M32
        good = tuple(w & m for w, m in zip(struct.unpack(">4I", ks)[:nw], tm))
        wrong = (good[0] ^ (1 << rnd.randrange(32)),) + good[1:]      # word 0 is always fully known
        words = np.array([struct.unpack(">4I", key), struct.unpack(">4I", iv)], np.int64)
        for tgt, want in ((good, 0), (wrong, -1)):
            ref = ad._k2_match(ad._k2_ik(key), iv, tgt, tuple(tm))
            got = int(scan(words, np.array([0]), np.array([1]), np.array([1]), tgt, tm)[0])
            bad += (got == 0) != ref or got != want
    return bad