package com.foxx.androidcast.network;

/*********************************************************************
 * UdpCastTransport.java
 * Created at: Fri 15 May 2026 16:49:02 +0200
 * Updated at: Wed 20 May 2026 15:17:13 +0200 by Anton Afanasyeu <a.afanasieff@gmail.com>
 * Commit: 5d8e82d2e60a21fff3138d2a394ee4e8b4c6dcb8
 * class UdpCastTransport
 * Contributors:
 *    - Anton Afanasyeu <a.afanasieff@gmail.com> (9 commits, 603 lines)
 *    - Cursor Agent (project assistant)
 * Digest: SHA256	4a5d534472ed2b021d783c8cf74246142994197c4c9c81c829154584fb6e6d9c
**********************************************************************/
import android.util.Log;

import com.foxx.androidcast.CastConfig;
import com.foxx.androidcast.network.transport.CombinedProtectionEngine;
import com.foxx.androidcast.network.transport.FecProtectionEngine;
import com.foxx.androidcast.network.transport.NackProtectionEngine;
import com.foxx.androidcast.network.transport.RetransmitCache;
import com.foxx.androidcast.network.transport.StreamProtectionEngine;
import com.foxx.androidcast.network.transport.StreamProtectionCompat;
import com.foxx.androidcast.network.transport.StreamProtectionMode;
import com.foxx.androidcast.network.transport.StreamProtectionScopes;
import com.foxx.androidcast.network.transport.TransportLossStats;
import com.foxx.androidcast.network.transport.TransportStatsProvider;
import com.foxx.androidcast.network.transport.fec.FecEncodedGroup;
import com.foxx.androidcast.network.transport.fec.FecShardWire;
import com.foxx.androidcast.stats.SessionStatsContext;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.SocketException;
import java.net.SocketTimeoutException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/** Best-effort message transport over UDP with fragmentation. */
public class UdpCastTransport implements CastTransport, TransportStatsProvider {
    private static final String TAG = "UdpCastTransport";
    private static final byte[] MAGIC = {'A', 'C', 'U', 'D'};
    /** magic(4) + type(1) + msgId(4) + fragIdx(2) + fragCnt(2) + totalLen(4) */
    private static final int HEADER = 17;
    private static final int MAX_RELIABLE_BYTES = 8 * 1024;

    private DatagramSocket socket;
    private InetSocketAddress peer;
    private int nextMessageId = 1;
    private volatile boolean listening;

    private final Map<Integer, Reassembly> pending = new ConcurrentHashMap<>();
    private final TransportLossStats lossStats = new TransportLossStats();
    private StreamProtectionEngine protectionEngine = StreamProtectionEngine.noop();
    private FecProtectionEngine fecEngine;
    private NackProtectionEngine nackEngine;
    private final RetransmitCache retransmitCache = new RetransmitCache();
    private CastFramingContext framingContext;
    private CastSessionGate inboundSessionGate;
    private static final long REASSEMBLY_TIMEOUT_MIN_MS = 3_000;
    private static final long REASSEMBLY_TIMEOUT_MAX_MS = 20_000;
    /** Large VP9 keyframes need many UDP fragments per FEC shard; allow extra time to collect shards. */
    private static final long FEC_GROUP_TIMEOUT_MS = 12_000;
    private static final int REASSEMBLY_LOG_EVERY = 16;
    private int reassemblyExpireLogCount;

    public boolean isListening() {
        return listening && socket != null && !socket.isClosed();
    }

    @Override
    public void prepareListen(int port) throws IOException {
        if (isListening()) {
            return;
        }
        openSocket();
        socket.bind(new InetSocketAddress(InetAddress.getByName("0.0.0.0"), port));
        socket.setSoTimeout(500);
        listening = true;
    }

    @Override
    public void listen(int port) throws IOException {
        prepareListen(port);
    }

    @Override
    public void connect(String host, int port) throws IOException {
        openSocket();
        peer = new InetSocketAddress(host, port);
        socket.setSoTimeout(500);
        listening = false;
    }

    private void openSocket() throws SocketException {
        if (socket != null && !socket.isClosed()) {
            return;
        }
        socket = new DatagramSocket(null);
        socket.setReuseAddress(true);
        socket.setReceiveBufferSize(2 * 1024 * 1024);
        socket.setSendBufferSize(2 * 1024 * 1024);
    }

    @Override
    public void setFramingContext(CastFramingContext context) {
        this.framingContext = context;
    }

    @Override
    public void setInboundSessionGate(CastSessionGate gate) {
        this.inboundSessionGate = gate;
    }

    public void applyStreamProtection(com.foxx.androidcast.CastSettings settings) {
        StreamProtectionEngine engine = com.foxx.androidcast.network.transport.StreamProtectionFactory
                .create(settings != null ? settings.getStreamProtection() : null, lossStats);
        setProtectionEngine(engine);
    }

    /** Select stream protection engine (FEC / NACK). */
    public void setProtectionEngine(StreamProtectionEngine engine) {
        protectionEngine = engine != null ? engine : StreamProtectionEngine.noop();
        fecEngine = FecProtectionEngine.extract(protectionEngine);
        nackEngine = extractNackEngine(protectionEngine);
        if (nackEngine != null) {
            nackEngine.setNackSender(messageId -> sendRaw(CastProtocol.MSG_NACK,
                    NackProtectionEngine.nackPayload(messageId)));
        }
        synchronized (lossStats) {
            lossStats.protectionMode = protectionEngine.mode();
        }
    }

    private static NackProtectionEngine extractNackEngine(StreamProtectionEngine engine) {
        if (engine instanceof NackProtectionEngine) {
            return (NackProtectionEngine) engine;
        }
        if (engine instanceof CombinedProtectionEngine) {
            return ((CombinedProtectionEngine) engine).nack();
        }
        return null;
    }

    @Override
    public TransportLossStats getLossStats() {
        synchronized (lossStats) {
            lossStats.publishLegacyCounters();
            return lossStats.copy();
        }
    }

    /** Send a low-level UDP control datagram (not FEC-wrapped). */
    public void sendRaw(byte type, byte[] payload) throws IOException {
        if (peer == null) {
            throw new IOException("UDP peer not set");
        }
        byte[] body = payload == null ? new byte[0] : payload;
        sendOnce(type, body, false);
    }

    void retransmitCached(byte wireType, byte[] body, int messageId) {
        try {
            sendOnceWithId(wireType, body, false, messageId);
        } catch (IOException ignored) {
        }
    }

    private void handleIncomingNack(int messageId) {
        protectionEngine.onNack(messageId);
        RetransmitCache.CachedMessage entry = retransmitCache.get(messageId);
        if (entry == null) {
            synchronized (lossStats) {
                lossStats.nackCacheMisses++;
            }
            if ((++reassemblyExpireLogCount % REASSEMBLY_LOG_EVERY) == 1) {
                Log.w(TAG, "NACK cache miss id=" + messageId);
            }
            return;
        }
        synchronized (lossStats) {
            lossStats.nackRetransmits++;
        }
        retransmitCached(entry.wireType, entry.body, messageId);
    }

    @Override
    public void send(byte type, byte[] payload) throws IOException {
        if (peer == null) {
            throw new IOException("UDP peer not set");
        }
        byte[] body;
        byte wireType = type;
        if (framingContext != null) {
            CastProtoHeader hdr = framingContext.newHeader(CastProtoHeader.HS_DEFAULT);
            body = CastPacketFramer.frame(hdr, type, payload);
            wireType = CastProtocol.MSG_WIRE_PACKET;
        } else {
            body = payload == null ? new byte[0] : payload;
        }
        if (fecEngine != null && StreamProtectionScopes.protectsUdpPayload(type)) {
            sendFecShardGroup(wireType, body);
            return;
        }
        body = protectionEngine.onOutboundPayload(type, body);
        if (body.length <= MAX_RELIABLE_BYTES) {
            sendOnce(wireType, body, true);
            sendOnce(wireType, body, true);
        } else {
            sendOnce(wireType, body, true);
        }
    }

    private void sendFecShardGroup(byte wireType, byte[] payload) throws IOException {
        FecEncodedGroup group = fecEngine.encodeGroup(payload);
        int baseId = nextMessageId;
        nextMessageId += group.totalShards();
        boolean cache = nackEngine != null;
        for (int i = 0; i < group.totalShards(); i++) {
            byte[] shardWire = fecEngine.encodeShardWire(group, baseId, i);
            sendOnceWithId(wireType, shardWire, cache, baseId + i);
            if (i + 1 < group.totalShards() && shardWire.length > 8_000) {
                paceBetweenLargeShards();
            }
        }
    }

    private static void paceBetweenLargeShards() {
        try {
            Thread.sleep(2);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }

    private void sendOnce(byte type, byte[] body, boolean cacheForNack) throws IOException {
        sendOnceWithId(type, body, cacheForNack, nextMessageId++);
    }

    private void sendOnceWithId(byte type, byte[] body, boolean cacheForNack, int messageId)
            throws IOException {
        if (cacheForNack) {
            retransmitCache.put(messageId, type, body);
        }
        int chunkSize = CastConfig.UDP_CHUNK_SIZE - HEADER;
        int fragments = Math.max(1, (body.length + chunkSize - 1) / chunkSize);
        boolean paceFragments = fragments > 24;
        for (int i = 0; i < fragments; i++) {
            int offset = i * chunkSize;
            int len = Math.min(chunkSize, body.length - offset);
            byte[] packet = buildPacket(type, messageId, i, fragments, body, offset, len, body.length);
            DatagramPacket dp = new DatagramPacket(packet, packet.length, peer);
            socket.send(dp);
            SessionStatsContext.onDatagramSent(packet.length);
            synchronized (lossStats) {
                lossStats.datagramsSent++;
            }
            if (paceFragments && i + 1 < fragments && (i & 7) == 7) {
                paceBetweenLargeShards();
            }
        }
    }

    @Override
    public CastProtocol.Message receive(int timeoutMs) throws IOException {
        if (socket == null || socket.isClosed()) {
            return null;
        }
        long deadline = System.currentTimeMillis() + timeoutMs;
        while (System.currentTimeMillis() < deadline) {
            expireStaleReassembly();
            expireStaleFecGroups();
            int remaining = (int) Math.max(1, deadline - System.currentTimeMillis());
            socket.setSoTimeout(remaining);
            try {
                byte[] buf = new byte[CastConfig.UDP_CHUNK_SIZE + 64];
                DatagramPacket dp = new DatagramPacket(buf, buf.length);
                socket.receive(dp);
                SessionStatsContext.onDatagramReceived(dp.getLength());
                synchronized (lossStats) {
                    lossStats.datagramsReceived++;
                }
                if (peer == null) {
                    peer = new InetSocketAddress(dp.getAddress(), dp.getPort());
                }
                CastProtocol.Message msg = parsePacket(dp.getData(), dp.getLength());
                if (msg != null) {
                    return msg;
                }
            } catch (SocketTimeoutException ignored) {
            }
        }
        return null;
    }

    private void expireStaleFecGroups() {
        if (fecEngine == null || nackEngine == null) {
            if (fecEngine != null) {
                fecEngine.expireStaleGroups(FEC_GROUP_TIMEOUT_MS, null);
            }
            return;
        }
        fecEngine.expireStaleGroups(FEC_GROUP_TIMEOUT_MS, nackEngine::requestNack);
    }

    private void expireStaleReassembly() {
        long now = System.currentTimeMillis();
        for (Map.Entry<Integer, Reassembly> e : pending.entrySet()) {
            Reassembly r = e.getValue();
            if (r != null && now > r.deadlineMs) {
                pending.remove(e.getKey());
                synchronized (lossStats) {
                    lossStats.reassemblyIncomplete++;
                }
                maybeLogReassemblyExpired(e.getKey(), r);
                int messageId = e.getKey();
                if (nackEngine != null) {
                    nackEngine.requestNack(messageId);
                } else if ((++reassemblyExpireLogCount % REASSEMBLY_LOG_EVERY) == 2) {
                    Log.w(TAG, "Reassembly expired without NACK id=" + messageId);
                }
            }
        }
    }

    private void maybeLogReassemblyExpired(int messageId, Reassembly r) {
        if ((++reassemblyExpireLogCount % REASSEMBLY_LOG_EVERY) != 1) {
            return;
        }
        Log.w(TAG, "UDP reassembly timeout id=" + messageId + " frags=" + r.received + "/"
                + r.fragCount + " bytes=" + r.totalLen);
    }

    private static long reassemblyDeadlineMs(int fragCount, int totalLen) {
        if (fragCount <= 1) {
            return System.currentTimeMillis() + REASSEMBLY_TIMEOUT_MIN_MS;
        }
        int chunkSize = Math.max(1, CastConfig.UDP_CHUNK_SIZE - HEADER);
        int expectedFrags = Math.max(fragCount, (totalLen + chunkSize - 1) / chunkSize);
        long budget = REASSEMBLY_TIMEOUT_MIN_MS + expectedFrags * 200L + (totalLen / 2048);
        if (fragCount > 1) {
            budget = Math.max(8_000L, budget);
        }
        budget = Math.min(REASSEMBLY_TIMEOUT_MAX_MS, budget);
        return System.currentTimeMillis() + budget;
    }

    private CastProtocol.Message parsePacket(byte[] data, int length) throws IOException {
        if (length < HEADER) {
            synchronized (lossStats) {
                lossStats.datagramsInvalid++;
            }
            return null;
        }
        for (int i = 0; i < 4; i++) {
            if (data[i] != MAGIC[i]) {
                synchronized (lossStats) {
                    lossStats.datagramsInvalid++;
                }
                return null;
            }
        }
        synchronized (lossStats) {
            lossStats.fragmentsReceived++;
        }
        byte type = data[4];
        int messageId = readInt(data, 5);
        int fragIndex = readShort(data, 9);
        int fragCount = readShort(data, 11);
        if (type == CastProtocol.MSG_NACK && fragCount <= 1) {
            if (nackEngine == null) {
                return null;
            }
            int payloadLenNack = length - HEADER;
            byte[] nackBody = payloadLenNack > 0 ? new byte[payloadLenNack] : new byte[0];
            if (payloadLenNack > 0) {
                System.arraycopy(data, HEADER, nackBody, 0, payloadLenNack);
            }
            int nackId = NackProtectionEngine.parseNackPayload(nackBody);
            if (nackId >= 0) {
                handleIncomingNack(nackId);
            }
            return null;
        }
        int totalLen = readInt(data, 13);
        int payloadLen = length - HEADER;
        byte[] chunk = new byte[payloadLen];
        System.arraycopy(data, HEADER, chunk, 0, payloadLen);

        CastProtocol.Message raw;
        if (fragCount <= 1) {
            byte[] body;
            if (totalLen != payloadLen && totalLen > 0 && totalLen <= payloadLen) {
                body = new byte[totalLen];
                System.arraycopy(chunk, 0, body, 0, totalLen);
            } else {
                body = chunk;
            }
            body = decodeInboundBody(type, body, messageId);
            if (body == null) {
                return null;
            }
            raw = new CastProtocol.Message(type, body);
            return unwrapWire(raw);
        }

        Reassembly existing = pending.get(messageId);
        if (existing != null && (existing.fragCount != fragCount || existing.totalLen != totalLen)) {
            pending.remove(messageId);
            synchronized (lossStats) {
                lossStats.reassemblySuperseded++;
            }
            existing = null;
        }
        Reassembly r = existing;
        if (r == null) {
            r = new Reassembly(fragCount, totalLen);
            pending.put(messageId, r);
        }
        if (!r.put(fragIndex, chunk, payloadLen)) {
            synchronized (lossStats) {
                lossStats.fragmentsDuplicate++;
            }
            return null;
        }
        if (!r.complete()) {
            return null;
        }
        pending.remove(messageId);
        synchronized (lossStats) {
            lossStats.reassemblyCompleted++;
        }
        byte[] assembled = r.assemble();
        assembled = decodeInboundBody(type, assembled, messageId);
        if (assembled == null) {
            return null;
        }
        return unwrapWire(new CastProtocol.Message(type, assembled));
    }

    private byte[] decodeInboundBody(byte type, byte[] body, int messageId) {
        StreamProtectionMode active = protectionEngine.mode();
        if (StreamProtectionCompat.shouldIgnoreInbound(body, active)) {
            synchronized (lossStats) {
                lossStats.unsupportedProtectionDropped++;
            }
            return null;
        }
        if (fecEngine != null && FecShardWire.isShardPacket(body)) {
            FecProtectionEngine.ShardDecodeResult shard = fecEngine.onInboundShard(type, body);
            if (shard.status == FecProtectionEngine.ShardDecodeResult.Status.PENDING) {
                return null;
            }
            if (shard.status == FecProtectionEngine.ShardDecodeResult.Status.FAILED) {
                if (nackEngine != null) {
                    nackEngine.requestNack(messageId);
                }
                return null;
            }
            return shard.payload;
        }
        byte[] decoded = protectionEngine.onInboundPayload(type, body);
        if (decoded != null
                && StreamProtectionCompat.shouldIgnoreInbound(decoded, active)) {
            synchronized (lossStats) {
                lossStats.unsupportedProtectionDropped++;
            }
            return null;
        }
        if (decoded == null && nackEngine != null && active != StreamProtectionMode.NONE) {
            nackEngine.requestNack(messageId);
        }
        return decoded;
    }

    private CastProtocol.Message unwrapWire(CastProtocol.Message raw) throws IOException {
        if (raw.type == CastProtocol.MSG_WIRE_PACKET) {
            if (inboundSessionGate != null) {
                CastProtoHeader hdr = CastPacketFramer.readHeader(raw.payload);
                if (!inboundSessionGate.accept(hdr)) {
                    synchronized (lossStats) {
                        lossStats.wirePacketsDropped++;
                    }
                    return null;
                }
            }
            try {
                return CastPacketFramer.unframe(raw.payload);
            } catch (IOException e) {
                synchronized (lossStats) {
                    lossStats.malformedPackets++;
                }
                throw e;
            }
        }
        return raw;
    }

    private static byte[] buildPacket(
            byte type, int messageId, int fragIndex, int fragCount,
            byte[] body, int offset, int len, int totalLen) {
        byte[] packet = new byte[HEADER + len];
        System.arraycopy(MAGIC, 0, packet, 0, 4);
        packet[4] = type;
        writeInt(packet, 5, messageId);
        writeShort(packet, 9, fragIndex);
        writeShort(packet, 11, fragCount);
        writeInt(packet, 13, totalLen);
        System.arraycopy(body, offset, packet, HEADER, len);
        return packet;
    }

    private static int readInt(byte[] b, int off) {
        return ((b[off] & 0xFF) << 24) | ((b[off + 1] & 0xFF) << 16)
                | ((b[off + 2] & 0xFF) << 8) | (b[off + 3] & 0xFF);
    }

    private static int readShort(byte[] b, int off) {
        return ((b[off] & 0xFF) << 8) | (b[off + 1] & 0xFF);
    }

    private static void writeInt(byte[] b, int off, int v) {
        b[off] = (byte) (v >> 24);
        b[off + 1] = (byte) (v >> 16);
        b[off + 2] = (byte) (v >> 8);
        b[off + 3] = (byte) v;
    }

    private static void writeShort(byte[] b, int off, int v) {
        b[off] = (byte) (v >> 8);
        b[off + 1] = (byte) v;
    }

    @Override
    public String getModeLabel() {
        return "UDP";
    }

    @Override
    public void releaseConnection() {
        peer = null;
        pending.clear();
    }

    @Override
    public void close() {
        listening = false;
        if (socket != null) {
            socket.close();
            socket = null;
        }
        pending.clear();
        peer = null;
    }

    private static final class Reassembly {
        final int fragCount;
        final int totalLen;
        final long startedMs;
        final long deadlineMs;
        final byte[][] parts;
        final int[] sizes;
        int received;

        Reassembly(int fragCount, int totalLen) {
            this.fragCount = fragCount;
            this.totalLen = totalLen;
            this.startedMs = System.currentTimeMillis();
            this.deadlineMs = reassemblyDeadlineMs(fragCount, totalLen);
            this.parts = new byte[fragCount][];
            this.sizes = new int[fragCount];
        }

        boolean put(int index, byte[] data, int len) {
            if (index < 0 || index >= fragCount || parts[index] != null) {
                return false;
            }
            byte[] copy = new byte[len];
            System.arraycopy(data, 0, copy, 0, len);
            parts[index] = copy;
            sizes[index] = len;
            received++;
            return true;
        }

        boolean complete() {
            return received == fragCount;
        }

        byte[] assemble() {
            ByteArrayOutputStream bos = new ByteArrayOutputStream(totalLen);
            for (int i = 0; i < fragCount; i++) {
                bos.write(parts[i], 0, sizes[i]);
            }
            byte[] out = bos.toByteArray();
            if (out.length > totalLen) {
                byte[] trimmed = new byte[totalLen];
                System.arraycopy(out, 0, trimmed, 0, totalLen);
                return trimmed;
            }
            return out;
        }
    }
}
