package com.foxx.androidcast.network.transport;

/*********************************************************************
 * FecProtectionEngine.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 FecProtectionEngine
 * Contributors:
 *    - Anton Afanasyeu <a.afanasieff@gmail.com> (4 commits, 277 lines)
 *    - Cursor Agent (project assistant)
 * Digest: SHA256	8440267db9d075ca98198af006689f0eb1aa6e5babddb197e421b39db055909e
**********************************************************************/
import com.foxx.androidcast.network.CastProtoHeader;
import com.foxx.androidcast.network.CastProtocol;
import com.foxx.androidcast.network.transport.fec.FecEncodedGroup;
import com.foxx.androidcast.network.transport.fec.FecShardReassembly;
import com.foxx.androidcast.network.transport.fec.FecShardWire;
import com.foxx.androidcast.network.transport.fec.ProtectionEnvelope;
import com.foxx.androidcast.network.transport.fec.ReedSolomonFec;

import android.util.Log;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/** Reed-Solomon FEC: one shard per UDP message; v1 monolithic decode still accepted inbound. */
public final class FecProtectionEngine implements StreamProtectionEngine {
    private static final String TAG = "FecProtection";
    private static final AtomicInteger NEXT_GROUP_ID = new AtomicInteger(1);
    private static final int LOG_EVERY_N = 32;
    private int inboundShardRejectLogCount;
    private int inboundDecodeFailLogCount;

    private final byte mode;
    private final int dataShards;
    private final int parityShards;
    private final TransportLossStats stats;
    private final StreamProtectionMode protectionMode;
    private final Map<Integer, FecShardReassembly> inboundGroups = new ConcurrentHashMap<>();

    public FecProtectionEngine(byte envelopeMode, int dataShards, int parityShards,
            TransportLossStats stats) {
        this.mode = envelopeMode;
        this.dataShards = dataShards;
        this.parityShards = parityShards;
        this.stats = stats;
        this.protectionMode = envelopeMode == ProtectionEnvelope.MODE_REED_SOLOMON
                ? StreamProtectionMode.FEC_REED_SOLOMON
                : StreamProtectionMode.FEC_3_4;
    }

    public static FecProtectionEngine extract(StreamProtectionEngine engine) {
        if (engine instanceof FecProtectionEngine) {
            return (FecProtectionEngine) engine;
        }
        if (engine instanceof CombinedProtectionEngine) {
            return ((CombinedProtectionEngine) engine).fec();
        }
        return null;
    }

    @Override
    public StreamProtectionMode mode() {
        return protectionMode;
    }

    /** Pick shard counts so each ACUD logical message stays small on lossy Wi‑Fi. */
    static int[] shardPlanForPayload(int payloadLength, int defaultData, int defaultParity) {
        if (payloadLength > 64_000) {
            return new int[] {16, 8};
        }
        if (payloadLength > 16_000) {
            return new int[] {8, 4};
        }
        return new int[] {defaultData, defaultParity};
    }

    public FecEncodedGroup encodeGroup(byte[] payload) throws IOException {
        if (payload == null || payload.length == 0) {
            throw new IOException("empty payload");
        }
        int[] plan = shardPlanForPayload(payload.length, dataShards, parityShards);
        int k = plan[0];
        int m = plan[1];
        ReedSolomonFec fec = new ReedSolomonFec(k, m);
        byte[][] shards = fec.encode(payload);
        int groupId = NEXT_GROUP_ID.getAndIncrement();
        if (stats != null) {
            synchronized (stats) {
                stats.fecPacketsEncoded += shards.length;
            }
        }
        return new FecEncodedGroup(mode, groupId, payload.length, k, m, shards);
    }

    public byte[] encodeShardWire(FecEncodedGroup group, int baseMessageId, int shardIndex)
            throws IOException {
        return FecShardWire.encodeShard(
                group.mode,
                group.groupId,
                baseMessageId,
                shardIndex,
                group.totalShards(),
                group.originalLen,
                group.dataShards,
                group.parityShards,
                group.shards[shardIndex]);
    }

  public static final class ShardDecodeResult {
        public enum Status { PENDING, OK, FAILED }

        public final Status status;
        public final byte[] payload;

        ShardDecodeResult(Status status, byte[] payload) {
            this.status = status;
            this.payload = payload;
        }
    }

    public ShardDecodeResult onInboundShard(byte type, byte[] payload) {
        if (!StreamProtectionScopes.protectsUdpPayload(type)) {
            return new ShardDecodeResult(ShardDecodeResult.Status.OK, payload);
        }
        if (type == CastProtocol.MSG_WIRE_PACKET && looksLikeWireHeader(payload)) {
            return new ShardDecodeResult(ShardDecodeResult.Status.OK, payload);
        }
        FecShardWire.Parsed shard = FecShardWire.tryParse(payload);
        if (shard == null) {
            byte[] legacy = onInboundPayload(type, payload);
            if (legacy == null) {
                return new ShardDecodeResult(ShardDecodeResult.Status.FAILED, null);
            }
            return new ShardDecodeResult(ShardDecodeResult.Status.OK, legacy);
        }
        FecShardReassembly assembly = inboundGroups.get(shard.groupId);
        if (assembly == null) {
            FecShardReassembly created = new FecShardReassembly(shard);
            FecShardReassembly prev = inboundGroups.putIfAbsent(shard.groupId, created);
            assembly = prev != null ? prev : created;
        }
        if (!assembly.addShard(shard)) {
            maybeLogShardReject(shard, assembly);
            byte[] retry = assembly.tryDecode();
            if (retry != null) {
                inboundGroups.remove(shard.groupId);
                if (stats != null) {
                    synchronized (stats) {
                        stats.fecPacketsDecoded++;
                    }
                }
                return new ShardDecodeResult(ShardDecodeResult.Status.OK, retry);
            }
            return new ShardDecodeResult(ShardDecodeResult.Status.PENDING, null);
        }
        byte[] decoded = assembly.tryDecode();
        if (decoded == null) {
            if (assembly.receivedCount() >= assembly.getTotalShards()) {
                if (stats != null) {
                    synchronized (stats) {
                        stats.fecDecodeFailures++;
                    }
                }
                maybeLogDecodeFail(shard.groupId, assembly.receivedCount(), assembly.getTotalShards());
                inboundGroups.remove(shard.groupId);
                return new ShardDecodeResult(ShardDecodeResult.Status.FAILED, null);
            }
            return new ShardDecodeResult(ShardDecodeResult.Status.PENDING, null);
        }
        inboundGroups.remove(shard.groupId);
        if (stats != null) {
            synchronized (stats) {
                stats.fecPacketsDecoded++;
            }
        }
        return new ShardDecodeResult(ShardDecodeResult.Status.OK, decoded);
    }

    public void expireStaleGroups(long maxAgeMs, ShardNackRequester nack) {
        long now = System.currentTimeMillis();
        Iterator<Map.Entry<Integer, FecShardReassembly>> it = inboundGroups.entrySet().iterator();
        while (it.hasNext()) {
            Map.Entry<Integer, FecShardReassembly> e = it.next();
            FecShardReassembly asm = e.getValue();
            if (now - asm.getStartedMs() <= maxAgeMs) {
                continue;
            }
            it.remove();
            if (nack != null) {
                for (int i = 0; i < asm.getTotalShards(); i++) {
                    if (!asm.hasShard(i)) {
                        nack.requestNack(asm.nackMessageIdForIndex(i));
                    }
                }
            }
            if (stats != null) {
                synchronized (stats) {
                    stats.fecDecodeFailures++;
                }
            }
        }
    }

    public interface ShardNackRequester {
        void requestNack(int messageId);
    }

    @Override
    public byte[] onOutboundPayload(byte type, byte[] payload) {
        return payload;
    }

    @Override
    public byte[] onInboundPayload(byte type, byte[] payload) {
        if (!StreamProtectionScopes.protectsUdpPayload(type)) {
            return payload;
        }
        if (type == CastProtocol.MSG_WIRE_PACKET && looksLikeWireHeader(payload)) {
            return payload;
        }
        if (FecShardWire.isShardPacket(payload)) {
            ShardDecodeResult shard = onInboundShard(type, payload);
            if (shard.status == ShardDecodeResult.Status.OK) {
                return shard.payload;
            }
            return null;
        }
        if (!ProtectionEnvelope.isProtected(payload)) {
            return payload;
        }
        ProtectionEnvelope.Decoded decoded = ProtectionEnvelope.tryDecode(payload);
        if (decoded == null) {
            return payload;
        }
        if (!decoded.complete) {
            if (stats != null) {
                synchronized (stats) {
                    stats.fecDecodeFailures++;
                }
            }
            return null;
        }
        if (stats != null) {
            synchronized (stats) {
                stats.fecPacketsDecoded++;
            }
        }
        return decoded.payload;
    }

    private void maybeLogShardReject(FecShardWire.Parsed shard, FecShardReassembly assembly) {
        if ((++inboundShardRejectLogCount % LOG_EVERY_N) != 1) {
            return;
        }
        Log.w(TAG, "FEC shard reject gid=" + shard.groupId + " idx=" + shard.shardIndex
                + " shardLen=" + shard.shardData.length + " recv=" + assembly.receivedCount()
                + "/" + assembly.getTotalShards());
    }

    private void maybeLogDecodeFail(int groupId, int received, int total) {
        if ((++inboundDecodeFailLogCount % LOG_EVERY_N) != 1) {
            return;
        }
        Log.w(TAG, "FEC RS decode failed gid=" + groupId + " shards=" + received + "/" + total);
    }

    private static boolean looksLikeWireHeader(byte[] payload) {
        if (payload == null || payload.length < 4) {
            return false;
        }
        int magic = ByteBuffer.wrap(payload, 0, 4).order(ByteOrder.LITTLE_ENDIAN).getInt();
        return magic == CastProtoHeader.MAGIC;
    }
}
