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

/*********************************************************************
 * ProtectionEnvelope.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 ProtectionEnvelope
 * Contributors:
 *    - Anton Afanasyeu <a.afanasieff@gmail.com> (2 commits, 158 lines)
 *    - Cursor Agent (project assistant)
 * Digest: SHA256	c6a05fbd2c4892e6995b46682b4aecf857deb9db651efada97c9cd86445ce2ee
**********************************************************************/
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;

/** Serialized FEC-protected UDP message body (before ACUD fragmentation). */
public final class ProtectionEnvelope {
    private static final byte[] MAGIC = {'A', 'C', 'P', 'F'};
    public static final byte VERSION = 1;
    private static final int HEADER_BYTES = 9;
    private static final int MAX_ORIGINAL_LENGTH = 16 * 1024 * 1024;

    public static final byte MODE_FEC_34 = 1;
    public static final byte MODE_REED_SOLOMON = 2;
    public static final byte MODE_FEC_NACK = 3;

    private ProtectionEnvelope() {}

    public static byte[] encode(byte mode, int dataShards, int parityShards, byte[] payload) throws IOException {
        ReedSolomonFec fec = new ReedSolomonFec(dataShards, parityShards);
        byte[][] shards = fec.encode(payload);
        ByteArrayOutputStream bos = new ByteArrayOutputStream();
        DataOutputStream dos = new DataOutputStream(bos);
        dos.write(MAGIC);
        dos.writeByte(VERSION);
        dos.writeByte(mode);
        dos.writeInt(payload.length);
        dos.writeByte(dataShards);
        dos.writeByte(parityShards);
        dos.writeByte(shards[0].length);
        for (byte[] shard : shards) {
            dos.write(shard);
        }
        return bos.toByteArray();
    }

    public static Decoded tryDecode(byte[] wire) {
        if (wire == null || wire.length < MAGIC.length + 10) {
            return null;
        }
        for (int i = 0; i < MAGIC.length; i++) {
            if (wire[i] != MAGIC[i]) {
                return null;
            }
        }
        try {
            DataInputStream dis = new DataInputStream(new ByteArrayInputStream(wire, MAGIC.length,
                    wire.length - MAGIC.length));
            byte ver = dis.readByte();
            if (ver != VERSION) {
                return null;
            }
            byte mode = dis.readByte();
            int originalLen = dis.readInt();
            int dataShards = dis.readUnsignedByte();
            int parityShards = dis.readUnsignedByte();
            int shardSize = dis.readUnsignedByte();
            if (!isValidHeader(dataShards, parityShards, shardSize, originalLen)) {
                return null;
            }
            int total = dataShards + parityShards;
            int expectedLen = MAGIC.length + HEADER_BYTES + total * shardSize;
            if (wire.length < expectedLen) {
                return null;
            }
            byte[][] shards = new byte[total][shardSize];
            for (int i = 0; i < total; i++) {
                dis.readFully(shards[i]);
            }
            ReedSolomonFec fec = new ReedSolomonFec(dataShards, parityShards);
            byte[] payload = fec.decode(shards, originalLen);
            if (payload == null) {
                return Decoded.incomplete(mode, dataShards, parityShards);
            }
            return Decoded.ok(mode, payload);
        } catch (IOException | RuntimeException e) {
            return null;
        }
    }

    public static boolean isProtected(byte[] wire) {
        if (wire == null || wire.length < MAGIC.length + HEADER_BYTES + 1) {
            return false;
        }
        for (int i = 0; i < MAGIC.length; i++) {
            if (wire[i] != MAGIC[i]) {
                return false;
            }
        }
        try {
            DataInputStream dis = new DataInputStream(new ByteArrayInputStream(wire, MAGIC.length,
                    wire.length - MAGIC.length));
            if (dis.readByte() != VERSION) {
                return false;
            }
            dis.readByte();
            int originalLen = dis.readInt();
            int dataShards = dis.readUnsignedByte();
            int parityShards = dis.readUnsignedByte();
            int shardSize = dis.readUnsignedByte();
            if (!isValidHeader(dataShards, parityShards, shardSize, originalLen)) {
                return false;
            }
            return wire.length >= MAGIC.length + HEADER_BYTES + (dataShards + parityShards) * shardSize;
        } catch (IOException e) {
            return false;
        }
    }

    private static boolean isValidHeader(int dataShards, int parityShards, int shardSize, int originalLen) {
        if (dataShards < 1 || parityShards < 1 || dataShards + parityShards > 255) {
            return false;
        }
        if (shardSize < 1) {
            return false;
        }
        if (originalLen < 1 || originalLen > MAX_ORIGINAL_LENGTH) {
            return false;
        }
        long maxRecoverable = (long) dataShards * shardSize;
        return originalLen <= maxRecoverable;
    }

    public static final class Decoded {
        public final boolean complete;
        public final byte mode;
        public final byte[] payload;
        public final int dataShards;
        public final int parityShards;

        private Decoded(boolean complete, byte mode, byte[] payload, int dataShards, int parityShards) {
            this.complete = complete;
            this.mode = mode;
            this.payload = payload;
            this.dataShards = dataShards;
            this.parityShards = parityShards;
        }

        static Decoded ok(byte mode, byte[] payload) {
            return new Decoded(true, mode, payload, 0, 0);
        }

        static Decoded incomplete(byte mode, int dataShards, int parityShards) {
            return new Decoded(false, mode, null, dataShards, parityShards);
        }
    }
}
