package com.foxx.androidcast.network.diag;

import android.util.Log;

import com.foxx.androidcast.network.CastProtocol;
import com.foxx.androidcast.network.CastSession;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Optional ~1 Hz diagnostic ping-pong during an active cast (developer setting).
 * Unknown {@link CastProtocol#MSG_DIAG_PING} on old peers is ignored by their handlers.
 */
public final class DiagPingSession {
    private static final String TAG = "DiagPing";
    private static final long PROBE_INTERVAL_MS = 1_000L;

    public interface SampleListener {
        void onDiagPingSample(DiagPingMetrics.Snapshot snapshot);
    }

    private final CastSession session;
    private final byte localResponseMode;
    private final DiagPingMetrics metrics = new DiagPingMetrics();
    private final ConcurrentHashMap<Long, Long> pendingProbeSentMs = new ConcurrentHashMap<>();
    private final SampleListener listener;

    private volatile boolean enabled;
    private long nextSeq;
    private long lastProbeMs;

    public DiagPingSession(CastSession session, byte localResponseMode, SampleListener listener) {
        this.session = session;
        this.localResponseMode = localResponseMode;
        this.listener = listener;
    }

    public void setEnabled(boolean on) {
        enabled = on;
        if (!on) {
            reset();
        }
    }

    public boolean isEnabled() {
        return enabled;
    }

    public DiagPingMetrics metrics() {
        return metrics;
    }

    public void reset() {
        nextSeq = 0;
        lastProbeMs = 0;
        pendingProbeSentMs.clear();
        metrics.reset();
    }

    public void tick(long nowMs) {
        if (!enabled) {
            return;
        }
        if (lastProbeMs == 0 || nowMs - lastProbeMs >= PROBE_INTERVAL_MS) {
            sendProbe(nowMs);
            lastProbeMs = nowMs;
        }
    }

    /** @return true if consumed (including invalid magic). */
    public boolean handleMessage(CastProtocol.Message msg) {
        if (msg.type != CastProtocol.MSG_DIAG_PING) {
            return false;
        }
        try {
            CastDiagPing.Packet p = CastProtocol.parseDiagPing(msg.payload);
            if (!p.isValidMagic()) {
                Log.w(TAG, "Ignoring diag ping with bad magic/version");
                return true;
            }
            if (p.isPong()) {
                onPong(p);
            } else {
                onProbe(p);
            }
            return true;
        } catch (IOException e) {
            Log.w(TAG, "Diag ping parse failed", e);
            return true;
        }
    }

    private void onProbe(CastDiagPing.Packet probe) throws IOException {
        CastDiagPing.Packet pong = new CastDiagPing.Packet();
        pong.flags = CastDiagPing.FLAG_PONG;
        pong.direction = probe.direction;
        pong.payload = probe.payload != null ? probe.payload : new byte[0];
        switch (probe.direction) {
            case CastDiagPing.DIR_MIRROR_AS_IS:
                pong.timestampMs = probe.timestampMs;
                pong.seq = probe.seq;
                break;
            case CastDiagPing.DIR_MIRROR_SMART:
                pong.timestampMs = System.currentTimeMillis();
                pong.seq = probe.seq;
                break;
            case CastDiagPing.DIR_REPLY_WITH:
                pong.timestampMs = System.currentTimeMillis();
                pong.seq = nextSeq++;
                pong.replyForSeq = probe.seq;
                pong.replyForTimestampMs = probe.timestampMs;
                break;
            default:
                pong.direction = CastDiagPing.DIR_MIRROR_SMART;
                pong.timestampMs = System.currentTimeMillis();
                pong.seq = probe.seq;
                break;
        }
        session.send(CastProtocol.MSG_DIAG_PING, CastProtocol.diagPingPayload(pong));
    }

    private void onPong(CastDiagPing.Packet pong) {
        long now = System.currentTimeMillis();
        long rtt;
        if (pong.direction == CastDiagPing.DIR_REPLY_WITH) {
            Long sentAt = pendingProbeSentMs.remove(pong.replyForSeq);
            rtt = sentAt != null ? now - sentAt : now - pong.replyForTimestampMs;
        } else {
            Long sentAt = pendingProbeSentMs.remove(pong.seq);
            rtt = sentAt != null ? now - sentAt : now - pong.timestampMs;
        }
        metrics.recordPong(rtt, pong.seq);
        notifyListener();
    }

    private void sendProbe(long nowMs) {
        try {
            long seq = nextSeq++;
            pendingProbeSentMs.put(seq, nowMs);
            metrics.recordProbeSent(seq);
            CastDiagPing.Packet probe = new CastDiagPing.Packet();
            probe.timestampMs = nowMs;
            probe.seq = seq;
            probe.direction = localResponseMode;
            probe.payload = new byte[0];
            session.send(CastProtocol.MSG_DIAG_PING, CastProtocol.diagPingPayload(probe));
        } catch (IOException e) {
            Log.w(TAG, "Send diag probe failed", e);
        }
    }

    private void notifyListener() {
        if (listener != null) {
            listener.onDiagPingSample(metrics.snapshot());
        }
    }
}
