package com.peachpasswords.codex

/**
 * Systematic Reed-Solomon erasure codec over GF(2⁸), Backblaze/JavaReedSolomon
 * construction (Vandermonde matrix made systematic). Mirrors the role of the
 * browser's `reed-solomon-erasure` WASM: encode fills parity shards in place;
 * reconstruct repairs any missing shards given an availability mask, tolerating
 * up to `parityShards` losses.
 *
 * Shards are equal-length byte arrays; the code operates byte-column-wise.
 */
object Galois {
    private val EXP = IntArray(512)
    private val LOG = IntArray(256)

    init {
        var x = 1
        for (i in 0 until 255) {
            EXP[i] = x
            LOG[x] = i
            x = x shl 1
            if (x and 0x100 != 0) x = x xor 0x11d
        }
        for (i in 255 until 512) EXP[i] = EXP[i - 255]
    }

    fun mul(a: Int, b: Int): Int =
        if (a == 0 || b == 0) 0 else EXP[LOG[a and 0xff] + LOG[b and 0xff]]

    fun div(a: Int, b: Int): Int {
        require(b != 0) { "divide by zero" }
        if (a == 0) return 0
        return EXP[(LOG[a and 0xff] - LOG[b and 0xff] + 255) % 255]
    }

    /** Field exponentiation a^n (used to build the Vandermonde matrix). */
    fun exp(a: Int, n: Int): Int {
        if (n == 0) return 1
        if (a == 0) return 0
        var result = 1
        repeat(n) { result = mul(result, a) }
        return result
    }
}

/** Dense GF(2⁸) matrix with the operations RS needs (multiply, invert). */
private class GfMatrix(val rows: Int, val cols: Int) {
    val data = Array(rows) { IntArray(cols) }

    fun times(other: GfMatrix): GfMatrix {
        require(cols == other.rows)
        val out = GfMatrix(rows, other.cols)
        for (r in 0 until rows) for (c in 0 until other.cols) {
            var v = 0
            for (k in 0 until cols) v = v xor Galois.mul(data[r][k], other.data[k][c])
            out.data[r][c] = v
        }
        return out
    }

    /** Gaussian-elimination inverse. Throws if singular. */
    fun invert(): GfMatrix {
        require(rows == cols)
        val work = GfMatrix(rows, cols * 2)
        for (r in 0 until rows) {
            System.arraycopy(data[r], 0, work.data[r], 0, cols)
            work.data[r][cols + r] = 1
        }
        work.gaussianEliminate()
        val out = GfMatrix(rows, cols)
        for (r in 0 until rows) System.arraycopy(work.data[r], cols, out.data[r], 0, cols)
        return out
    }

    private fun gaussianEliminate() {
        for (r in 0 until rows) {
            if (data[r][r] == 0) {
                var swap = r + 1
                while (swap < rows && data[swap][r] == 0) swap++
                require(swap < rows) { "matrix is singular" }
                val tmp = data[r]; data[r] = data[swap]; data[swap] = tmp
            }
            val pivot = data[r][r]
            if (pivot != 1) for (c in 0 until cols) data[r][c] = Galois.div(data[r][c], pivot)
            for (other in 0 until rows) {
                if (other == r) continue
                val factor = data[other][r]
                if (factor != 0) for (c in 0 until cols) {
                    data[other][c] = data[other][c] xor Galois.mul(factor, data[r][c])
                }
            }
        }
    }

    companion object {
        /** Encoding matrix: rows 0..dataShards-1 form the identity (systematic). */
        fun buildEncoding(dataShards: Int, totalShards: Int): GfMatrix {
            val vander = GfMatrix(totalShards, dataShards)
            for (r in 0 until totalShards) for (c in 0 until dataShards) {
                vander.data[r][c] = Galois.exp(r, c)
            }
            val top = GfMatrix(dataShards, dataShards)
            for (r in 0 until dataShards) System.arraycopy(vander.data[r], 0, top.data[r], 0, dataShards)
            return vander.times(top.invert())
        }
    }
}

object ReedSolomon {

    /** Fill the parity shards (indices dataShards..total-1) in place. */
    fun encode(shards: Array<ByteArray>, dataShards: Int, parityShards: Int) {
        val total = dataShards + parityShards
        require(shards.size == total) { "shard count mismatch" }
        val matrix = GfMatrix.buildEncoding(dataShards, total)
        val shardSize = shards[0].size
        for (parityRow in dataShards until total) {
            val out = shards[parityRow]
            for (b in 0 until shardSize) {
                var v = 0
                for (dataRow in 0 until dataShards) {
                    v = v xor Galois.mul(matrix.data[parityRow][dataRow], shards[dataRow][b].toInt() and 0xff)
                }
                out[b] = v.toByte()
            }
        }
    }

    /**
     * Reconstruct missing shards in place. [present] marks which shards were
     * recovered intact; at least [dataShards] must be present. Repairs the
     * data shards (and, if asked, parity) by inverting the sub-matrix of
     * present rows.
     */
    fun reconstruct(
        shards: Array<ByteArray>,
        dataShards: Int,
        parityShards: Int,
        present: BooleanArray,
    ) {
        val total = dataShards + parityShards
        val shardSize = shards.first { it.isNotEmpty() }.size
        val presentCount = present.count { it }
        if (presentCount < dataShards) {
            throw CodexException(CodexErrorCode.INSUFFICIENT_CHUNKS, "too few shards present: $presentCount < $dataShards")
        }

        val matrix = GfMatrix.buildEncoding(dataShards, total)

        // Build a dataShards×dataShards sub-matrix from the first dataShards
        // present rows, plus the corresponding present shard data.
        val sub = GfMatrix(dataShards, dataShards)
        val subShards = arrayOfNulls<ByteArray>(dataShards)
        var subRow = 0
        var r = 0
        while (subRow < dataShards) {
            if (present[r]) {
                System.arraycopy(matrix.data[r], 0, sub.data[subRow], 0, dataShards)
                subShards[subRow] = shards[r]
                subRow++
            }
            r++
        }
        val decode = sub.invert()

        // Recover every missing data shard.
        for (missing in 0 until dataShards) {
            if (present[missing]) continue
            val out = ByteArray(shardSize)
            for (b in 0 until shardSize) {
                var v = 0
                for (k in 0 until dataShards) {
                    v = v xor Galois.mul(decode.data[missing][k], subShards[k]!![b].toInt() and 0xff)
                }
                out[b] = v.toByte()
            }
            shards[missing] = out
            present[missing] = true
        }

        // Recompute any missing parity from the now-complete data shards.
        for (parityRow in dataShards until total) {
            if (present[parityRow]) continue
            val out = ByteArray(shardSize)
            for (b in 0 until shardSize) {
                var v = 0
                for (dataRow in 0 until dataShards) {
                    v = v xor Galois.mul(matrix.data[parityRow][dataRow], shards[dataRow][b].toInt() and 0xff)
                }
                out[b] = v.toByte()
            }
            shards[parityRow] = out
            present[parityRow] = true
        }
    }
}