import org.bouncycastle.jcajce.provider.digest.SHA3;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumKeyGenerationParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumKeyPairGenerator;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumPrivateKeyParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumPublicKeyParameters;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.SecureRandom;
import java.util.Arrays;

/**
 * Prints deterministic wallet vectors for JS parity checks.
 * Compile/run against bcprov-jdk18on on the classpath.
 */
public class ParityTool {
    private static final String ALPHABET = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz";

    public static void main(String[] args) throws Exception {
        // Fixed 64-byte Dilithium input seed (simulates Wallet.deriveKeyPair output).
        byte[] fullSeed = new byte[64];
        for (int i = 0; i < 64; i++) fullSeed[i] = (byte) i;

        byte[] sha3Seed = sha3_256(fullSeed);
        SecureRandom rng = SecureRandom.getInstance("SHA1PRNG");
        rng.setSeed(sha3Seed);
        byte[] dilSeed = new byte[32];
        rng.nextBytes(dilSeed);

        DilithiumKeyPairGenerator keyGen = new DilithiumKeyPairGenerator();
        SecureRandom keyRng = SecureRandom.getInstance("SHA1PRNG");
        keyRng.setSeed(sha3Seed);
        keyGen.init(new DilithiumKeyGenerationParameters(keyRng, DilithiumParameters.dilithium5));
        var pair = keyGen.generateKeyPair();
        byte[] pub = ((DilithiumPublicKeyParameters) pair.getPublic()).getEncoded();
        byte[] priv = ((DilithiumPrivateKeyParameters) pair.getPrivate()).getEncoded();
        String address = "SRMD" + base58(sha3_256(pub));

        System.out.println("fullSeed=" + hex(fullSeed));
        System.out.println("sha3Seed=" + hex(sha3Seed));
        System.out.println("sha1prng32=" + hex(dilSeed));
        System.out.println("pubLen=" + pub.length);
        System.out.println("privLen=" + priv.length);
        System.out.println("pub=" + hex(pub));
        System.out.println("address=" + address);

        // BIP39-style derive path with empty passphrase seed material:
        // SHA256(seed||"sarmad-wallet-0") || SHA256(seed||"sarmad-wallet-0-1")
        byte[] bipSeed = new byte[64];
        Arrays.fill(bipSeed, (byte) 7);
        MessageDigest sha256 = MessageDigest.getInstance("SHA-256");
        sha256.update(bipSeed);
        sha256.update("sarmad-wallet-0".getBytes(StandardCharsets.UTF_8));
        byte[] a = sha256.digest();
        sha256.reset();
        sha256.update(bipSeed);
        sha256.update("sarmad-wallet-0-1".getBytes(StandardCharsets.UTF_8));
        byte[] b = sha256.digest();
        byte[] derived = new byte[64];
        System.arraycopy(a, 0, derived, 0, 32);
        System.arraycopy(b, 0, derived, 32, 32);
        System.out.println("deriveExample=" + hex(derived));
    }

    private static byte[] sha3_256(byte[] in) {
        SHA3.Digest256 d = new SHA3.Digest256();
        return d.digest(in);
    }

    private static String hex(byte[] b) {
        StringBuilder sb = new StringBuilder(b.length * 2);
        for (byte x : b) sb.append(String.format("%02x", x));
        return sb.toString();
    }

    private static String base58(byte[] input) {
        if (input.length == 0) return "";
        byte[] bytes = Arrays.copyOf(input, input.length);
        int count = 0;
        for (byte b : bytes) {
            if (b == 0) count++;
            else break;
        }
        byte[] temp = new byte[bytes.length * 2];
        int j = temp.length;
        int startAt = count;
        while (startAt < bytes.length) {
            int carry = bytes[startAt] & 0xFF;
            int i = 0;
            int k = temp.length - 1;
            while (k >= 0 && (carry != 0 || i < j - (temp.length - bytes.length))) {
                carry += 256 * (temp[k] & 0xFF);
                temp[k] = (byte) (carry % 58);
                carry /= 58;
                k--;
                i++;
            }
            startAt++;
        }
        StringBuilder result = new StringBuilder();
        for (byte b : temp) {
            if (b != 0 || result.length() > 0) result.append(ALPHABET.charAt(b & 0xFF));
        }
        for (int i = 0; i < count; i++) result.insert(0, ALPHABET.charAt(0));
        return result.toString();
    }
}
