package com.foxx.androidcast.network.transport.fec;

/*********************************************************************
 * ReedSolomonFec.java
 * Created at: Sun 17 May 2026 22:49:49 +0200
 * Updated at: Wed 20 May 2026 15:17:13 +0200 by Anton Afanasyeu <a.afanasieff@gmail.com>
 * Commit: 5d8e82d2e60a21fff3138d2a394ee4e8b4c6dcb8
 * class ReedSolomonFec
 * Contributors:
 *    - Anton Afanasyeu <a.afanasieff@gmail.com> (2 commits, 182 lines)
 *    - Cursor Agent (project assistant)
 * Digest: SHA256	f6b72854c1477ec632367c3c6a88aa6592c7a1ff1fc3255c6ebb62c861a06031
**********************************************************************/
import java.util.Arrays;

/**
 * Reed-Solomon erasure code over GF(256). Requires any {@code dataShards} of
 * {@code dataShards + parityShards} shards to recover the payload.
 */
public final class ReedSolomonFec {
    private final int dataShards;
    private final int parityShards;
    private final int totalShards;
    private final int[] encodeMatrix;

    public ReedSolomonFec(int dataShards, int parityShards) {
        if (dataShards < 1 || parityShards < 1 || dataShards + parityShards > 255) {
            throw new IllegalArgumentException("Invalid shard counts");
        }
        this.dataShards = dataShards;
        this.parityShards = parityShards;
        this.totalShards = dataShards + parityShards;
        this.encodeMatrix = buildEncodeMatrix(dataShards, parityShards);
    }

    public int dataShards() {
        return dataShards;
    }

    public int parityShards() {
        return parityShards;
    }

    public int totalShards() {
        return totalShards;
    }

    public static int shardSizeFor(int payloadLength, int dataShards) {
        return Math.max(1, (payloadLength + dataShards - 1) / dataShards);
    }

    public byte[][] encode(byte[] payload) {
        int shardSize = shardSizeFor(payload.length, dataShards);
        byte[][] shards = new byte[totalShards][shardSize];
        for (int i = 0; i < dataShards; i++) {
            int off = i * shardSize;
            int len = Math.min(shardSize, Math.max(0, payload.length - off));
            if (len > 0) {
                System.arraycopy(payload, off, shards[i], 0, len);
            }
        }
        for (int p = 0; p < parityShards; p++) {
            for (int d = 0; d < dataShards; d++) {
                int coeff = encodeMatrix[p * dataShards + d] & 0xFF;
                GaloisField.mulAddInto(shards[dataShards + p], shards[d], coeff);
            }
        }
        return shards;
    }

    public byte[] decode(byte[][] shards, int originalLength) {
        if (shards == null || shards.length != totalShards) {
            return null;
        }
        int shardSize = 0;
        for (byte[] s : shards) {
            if (s != null && s.length > shardSize) {
                shardSize = s.length;
            }
        }
        if (shardSize == 0) {
            return null;
        }
        int[] idx = new int[dataShards];
        int found = 0;
        for (int i = 0; i < totalShards && found < dataShards; i++) {
            if (shards[i] != null) {
                idx[found++] = i;
            }
        }
        if (found < dataShards) {
            return null;
        }
        int[][] matrix = new int[dataShards][dataShards];
        for (int row = 0; row < dataShards; row++) {
            int shardIndex = idx[row];
            if (shardIndex < dataShards) {
                matrix[row][shardIndex] = 1;
            } else {
                int p = shardIndex - dataShards;
                for (int col = 0; col < dataShards; col++) {
                    matrix[row][col] = encodeMatrix[p * dataShards + col] & 0xFF;
                }
            }
        }
        int[][] inv = invertMatrix(matrix);
        if (inv == null) {
            return null;
        }
        byte[][] data = new byte[dataShards][shardSize];
        for (int row = 0; row < dataShards; row++) {
            Arrays.fill(data[row], (byte) 0);
            for (int col = 0; col < dataShards; col++) {
                int coeff = inv[row][col] & 0xFF;
                if (coeff != 0) {
                    GaloisField.mulAddInto(data[row], shards[idx[col]], coeff);
                }
            }
        }
        byte[] out = new byte[originalLength];
        int shardSizeOut = shardSizeFor(originalLength, dataShards);
        for (int i = 0; i < dataShards; i++) {
            int off = i * shardSizeOut;
            int len = Math.min(shardSizeOut, originalLength - off);
            if (len > 0) {
                len = Math.min(len, data[i].length);
                System.arraycopy(data[i], 0, out, off, len);
            }
        }
        return out;
    }

    private static int[] buildEncodeMatrix(int k, int m) {
        int[] matrix = new int[m * k];
        for (int row = 0; row < m; row++) {
            for (int col = 0; col < k; col++) {
                matrix[row * k + col] = GaloisField.pow((row + 1) & 0xFF, col);
            }
        }
        return matrix;
    }

    private static int[][] invertMatrix(int[][] matrix) {
        int n = matrix.length;
        int[][] aug = new int[n][2 * n];
        for (int i = 0; i < n; i++) {
            System.arraycopy(matrix[i], 0, aug[i], 0, n);
            aug[i][n + i] = 1;
        }
        for (int col = 0; col < n; col++) {
            int pivot = -1;
            for (int row = col; row < n; row++) {
                if (aug[row][col] != 0) {
                    pivot = row;
                    break;
                }
            }
            if (pivot < 0) {
                return null;
            }
            if (pivot != col) {
                int[] tmp = aug[pivot];
                aug[pivot] = aug[col];
                aug[col] = tmp;
            }
            int inv = GaloisField.div(1, aug[col][col]);
            for (int j = 0; j < 2 * n; j++) {
                aug[col][j] = GaloisField.mul(aug[col][j], inv);
            }
            for (int row = 0; row < n; row++) {
                if (row == col || aug[row][col] == 0) {
                    continue;
                }
                int factor = aug[row][col];
                for (int j = 0; j < 2 * n; j++) {
                    aug[row][j] = GaloisField.add(aug[row][j], GaloisField.mul(factor, aug[col][j]));
                }
            }
        }
        int[][] out = new int[n][n];
        for (int i = 0; i < n; i++) {
            System.arraycopy(aug[i], n, out[i], 0, n);
        }
        return out;
    }
}
