package org.bouncycastle.crypto.signers.mldsa;

class Rounding
{
    static void power2RoundAll(int[] c0, int[] c1)
    {
        int d = MLDSAEngine.DilithiumD, n = MLDSAEngine.DilithiumN;
        int u = (1 << (d - 1)) - 1, v = -1 << d;

        for (int i = 0; i < n; ++i)
        {
            int a = c0[i];

            int t = a + u;
            int r1 = a - (t & v);

            c0[i] = t >> d;
            c1[i] = r1;
        }
    }
    
    public static long decompose(int a, int gamma2)
    {
        int a1, a0;

        a1 = (a + 127) >> 7;
        if (gamma2 == (MLDSAEngine.DilithiumQ - 1) / 32)
        {
            a1 = (a1 * 1025 + (1 << 21)) >> 22;
            a1 &= 15;
        }
        else if (gamma2 == (MLDSAEngine.DilithiumQ - 1) / 88)
        {
            a1 = (a1 * 11275 + (1 << 23)) >> 24;
            a1 ^= ((43 - a1) >> 31) & a1;
        }
        else
        {
            throw new RuntimeException("Wrong Gamma2!");
        }

        a0 = a - a1 * 2 * gamma2;
        a0 -= (((MLDSAEngine.DilithiumQ - 1) / 2 - a0) >> 31) & MLDSAEngine.DilithiumQ;
        return ((long)a0 << 32) | (a1 & 0xFFFFFFFFL);
    }

    // Constant-time note: this branches on its inputs, but the returned value is the hint bit that
    // ships in the signature — information-equivalent to public output. Matches reference make_hint.
    public static int makeHint(int a0, int a1, MLDSAEngine engine)
    {
        int g2 = engine.getDilithiumGamma2(), q = MLDSAEngine.DilithiumQ;
        if (a0 <= g2 || a0 > q - g2 || (a0 == q - g2 && a1 == 0))
        {
            return 0;
        }
        return 1;
    }

    public static int useHint(int a, int hint, int gamma2)
    {
        long packed = decompose(a, gamma2);
        int a0 = (int)(packed >> 32);
        int a1 = (int)packed;

        if (hint == 0)
        {
            return a1;
        }

        if (gamma2 == (MLDSAEngine.DilithiumQ - 1) / 32)
        {
            if (a0 > 0)
            {
                return (a1 + 1) & 15;
            }
            else
            {
                return (a1 - 1) & 15;
            }
        }
        else if (gamma2 == (MLDSAEngine.DilithiumQ - 1) / 88)
        {
            if (a0 > 0)
            {
                return (a1 == 43) ? 0 : a1 + 1;
            }
            else
            {
                return (a1 == 0) ? 43 : a1 - 1;
            }
        }
        else
        {
            throw new RuntimeException("Wrong Gamma2!");
        }
    }
}
