package invalid.lena.scrcpy;

import android.media.MediaCodec;
import android.media.MediaFormat;
import android.os.Handler;
import android.os.HandlerThread;
import android.view.Surface;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayDeque;
import java.util.Deque;
import java.util.Iterator;

// MediaCodec async-mode video decoder writing to a Surface.
//
// MediaCodec hands us input buffer indices on its callback handler;
// VideoStream pushes encoded frames at us synchronously. The two ends
// meet through a small queue of pending frames waiting for input buffers,
// plus a corresponding pool of free input buffer indices.
//
// Back-pressure policy: when no input buffer is free and the pending
// queue is full, replace the oldest frame of the same kind, or the oldest
// frame overall. The queue remains bounded even if a peer floods CSD.
// Keyframes are not distinguishable here (the flag stays in VideoStream)
// so they drop like any delta frame; the picture heals at the next one.
//
// Output timing: releaseOutputBuffer is called with an absolute nano
// timestamp on the System.nanoTime clock, derived from the wire PTS so
// the surface composer interpolates frames against vsync instead of
// rendering them as fast as they decode. Anchor is set on the first
// frame: wallClockNs0 - ptsUs0 * 1000 = constant offset, then for each
// subsequent buffer renderAt = ptsUs * 1000 + offset.
public final class VideoSink implements VideoFrames {

    private static final int MAX_PENDING = 8;

    private volatile Surface surface;
    private final Runnable onFatalError;
    public  volatile long    frames; // public read for the status overlay
    private volatile MediaCodec codec;
    private HandlerThread handlerThread;
    private Handler handler;

    private final Object lock = new Object();
    private final Deque<Integer> freeInputs = new ArrayDeque<>(16);
    private final Deque<Frame>   pending    = new ArrayDeque<>(MAX_PENDING);

    private boolean released;
    private long    ptsOffsetNs;  // wall-ns = ptsUs * 1000 + ptsOffsetNs
    private boolean ptsAnchored;

    private static final class Frame {
        final byte[] data;
        final long   ptsUs;
        final boolean isConfig;
        Frame(byte[] d, long pts, boolean cfg) { data = d; ptsUs = pts; isConfig = cfg; }
    }

    public VideoSink(Surface surface, Runnable onFatalError) {
        this.surface = surface;
        this.onFatalError = onFatalError;
    }

    // Swap the output Surface without rebuilding MediaCodec. Passing null
    // detaches output: the codec keeps decoding but discards frames so the
    // wire stays drained while the activity is backgrounded. Passing a new
    // Surface re-attaches and rendering resumes from the next decoded frame.
    public void setOutputSurface(Surface newSurface) {
        synchronized (lock) {
            if (released) return;
            this.surface = newSurface;
            MediaCodec c = codec;
            if (c == null) return; // not yet configured; new surface will be used at configure time
            try {
                if (newSurface != null) c.setOutputSurface(newSurface);
                // Calling setOutputSurface(null) is unsupported on some
                // codecs; instead, ignore output buffers in the callback
                // when surface is null (see onOutputBufferAvailable).
            } catch (IllegalStateException e) {
                Log.w("video sink: setOutputSurface: %s", e);
            }
            // Re-anchor PTS so frames decoded against the old surface clock
            // don't drag the new surface's render time into the past.
            ptsAnchored = false;
        }
    }

    @Override
    public void configure(int codecFourcc, int width, int height) throws IOException {
        String mime = mimeFor(codecFourcc);
        if (mime == null) throw new IOException("unsupported video codec " + Wire.fourccName(codecFourcc));
        Log.i("video sink: configure mime=%s %dx%d", mime, width, height);

        handlerThread = new HandlerThread("video-mc");
        handlerThread.start();
        handler = new Handler(handlerThread.getLooper());

        codec = MediaCodec.createDecoderByType(mime);
        codec.setCallback(new MediaCodec.Callback() {
            @Override public void onInputBufferAvailable(MediaCodec mc, int idx) {
                onFreeInput(mc, idx);
            }
            @Override public void onOutputBufferAvailable(MediaCodec mc, int idx, MediaCodec.BufferInfo info) {
                try {
                    if (surface == null) {
                        // Activity backgrounded - discard output instead
                        // of rendering to a dead surface.
                        mc.releaseOutputBuffer(idx, false);
                        return;
                    }
                    // PTS-honoured render: SurfaceFlinger queues the buffer
                    // for renderTimestampNs and interpolates against vsync,
                    // so bursty arrivals smooth out instead of judder.
                    // Config frames (pts=0) fall through to render-immediately.
                    long renderAtNs = renderTimeNs(info.presentationTimeUs);
                    if (renderAtNs == 0L) mc.releaseOutputBuffer(idx, true);
                    else                  mc.releaseOutputBuffer(idx, renderAtNs);
                } catch (IllegalStateException ignored) {}
            }
            @Override public void onError(MediaCodec mc, MediaCodec.CodecException e) {
                Log.e(e, "video sink: codec error");
                if (mc == codec && onFatalError != null) onFatalError.run();
            }
            @Override public void onOutputFormatChanged(MediaCodec mc, MediaFormat fmt) {
                Log.i("video sink: output format %s", fmt);
            }
        }, handler);

        MediaFormat fmt = MediaFormat.createVideoFormat(mime, width, height);
        codec.configure(fmt, surface, null, 0);
        codec.start();
    }

    // Called by VideoStream for every encoded frame, in order.
    @Override
    public void feed(byte[] data, long ptsUs, boolean isConfig) {
        if (!isConfig) frames++;
        synchronized (lock) {
            if (released) return;
            // Try to drain immediately if there's a free input.
            while (!pending.isEmpty() && !freeInputs.isEmpty()) {
                submit(codec, pending.pollFirst(), freeInputs.pollFirst());
            }
            if (!freeInputs.isEmpty()) {
                submit(codec, new Frame(data, ptsUs, isConfig), freeInputs.pollFirst());
                return;
            }
            // Keep the queue strictly bounded. A newer config frame replaces
            // an older one; retaining every config packet lets a hostile peer
            // turn the queue into an unbounded allocation sink.
            if (pending.size() >= MAX_PENDING) {
                boolean removed = false;
                for (Iterator<Frame> it = pending.iterator(); it.hasNext(); ) {
                    Frame f = it.next();
                    if (f.isConfig == isConfig) {
                        it.remove();
                        removed = true;
                        break;
                    }
                }
                if (!removed) pending.pollFirst();
            }
            pending.offerLast(new Frame(data, ptsUs, isConfig));
        }
    }

    @Override
    public void release() {
        synchronized (lock) {
            if (released) return;
            released = true;
        }
        teardownCodec();
    }

    // Tear down the current decoder and reconfigure with new dimensions.
    // The next CSD frame on the wire (server resets on resize) will prime
    // the new codec instance.
    @Override
    public void reconfigure(int codecFourcc, int width, int height) throws IOException {
        synchronized (lock) {
            if (released) return;
        }
        teardownCodec();
        synchronized (lock) {
            freeInputs.clear();
            pending.clear();
            ptsAnchored = false; // re-anchor on the first frame of the new run
        }
        configure(codecFourcc, width, height);
    }

    private void teardownCodec() {
        MediaCodec c = codec;
        codec = null;
        HandlerThread ht = handlerThread;
        handlerThread = null;
        handler = null;
        if (c != null) {
            try { c.stop(); } catch (Exception ignored) {}
            try { c.release(); } catch (Exception ignored) {}
        }
        if (ht != null) ht.quitSafely();
    }

    // Internal - runs on the MediaCodec callback thread.
    private void onFreeInput(MediaCodec mc, int idx) {
        synchronized (lock) {
            if (released || mc != codec) return;
            if (!pending.isEmpty()) submit(mc, pending.pollFirst(), idx);
            else                    freeInputs.offerLast(idx);
        }
    }

    // Must be called with `lock` held. codec can be null mid-reconfigure
    // (teardownCodec runs unlocked); the frame is dropped like any other
    // back-pressure casualty.
    private void submit(MediaCodec mc, Frame f, int idx) {
        if (mc == null || mc != codec) return;
        try {
            ByteBuffer buf = mc.getInputBuffer(idx);
            if (buf == null || f.data.length > buf.capacity()) {
                Log.e("video sink: frame exceeds codec input (%d bytes)", f.data.length);
                if (onFatalError != null) onFatalError.run();
                return;
            }
            buf.clear();
            buf.put(f.data);
            int flags = f.isConfig ? MediaCodec.BUFFER_FLAG_CODEC_CONFIG : 0;
            mc.queueInputBuffer(idx, 0, f.data.length, f.ptsUs, flags);
        } catch (IllegalStateException e) {
            Log.w("video sink: queueInputBuffer: %s", e);
        }
    }

    // Convert a wire PTS (microseconds since some scrcpy epoch) into a
    // System.nanoTime value the surface composer should render at.
    // First call anchors the offset to "now" so latency stays whatever
    // the wire produced. PTS=0 (config frames) and unanchored state
    // both return 0 → caller falls back to render-immediately.
    private long renderTimeNs(long ptsUs) {
        if (ptsUs <= 0L) return 0L;
        synchronized (lock) {
            if (!ptsAnchored) {
                ptsOffsetNs = System.nanoTime() - ptsUs * 1000L;
                ptsAnchored = true;
            }
            return ptsUs * 1000L + ptsOffsetNs;
        }
    }

    private static String mimeFor(int fourcc) {
        switch (fourcc) {
            case Wire.CODEC_H264: return MediaFormat.MIMETYPE_VIDEO_AVC;
            case Wire.CODEC_H265: return MediaFormat.MIMETYPE_VIDEO_HEVC;
            case Wire.CODEC_AV1:  return MediaFormat.MIMETYPE_VIDEO_AV1;
            default: return null;
        }
    }
}