package com.foxx.androidcast.media.codec;

/*********************************************************************
 * LibvpxVideoEncoderSink.java
 * Created at: Mon 18 May 2026 14:25:02 +0200
 * Updated at: Wed 20 May 2026 15:17:13 +0200 by Anton Afanasyeu <a.afanasieff@gmail.com>
 * Commit: 5d8e82d2e60a21fff3138d2a394ee4e8b4c6dcb8
 * class LibvpxVideoEncoderSink
 * Contributors:
 *    - Anton Afanasyeu <a.afanasieff@gmail.com> (4 commits, 244 lines)
 *    - Cursor Agent (project assistant)
 * Digest: SHA256	43faed143aa87bcb23e317cbe1118de3e2e598bfcf5ad1d16edd0c1265b43e98
**********************************************************************/
import android.graphics.Bitmap;
import android.graphics.PixelFormat;
import android.media.Image;
import android.media.ImageReader;
import android.media.MediaFormat;
import android.os.Handler;
import android.os.HandlerThread;
import android.util.Log;
import android.view.Surface;

import com.foxx.androidcast.CastSettings;
import com.foxx.androidcast.CastTuningEngine;
import com.foxx.androidcast.media.RgbToYuv420;
import com.foxx.androidcast.media.codec.jni.NativeCodecBridge;
import com.foxx.androidcast.media.codec.jni.VpxEncodedFrame;
import com.foxx.androidcast.sender.VideoEncoder;

import java.io.IOException;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;

/**
 * Software VP8/VP9 via libvpx when {@link NativeCodecBridge#isLibvpxAvailable()}.
 * Screen capture uses an {@link ImageReader} surface; calibration may use buffer or GLES paths.
 */
public final class LibvpxVideoEncoderSink implements VideoEncoderSink {
    private static final String TAG = "LibvpxVideoEncoder";

    private final CastSettings.VideoCodec preference;
    private final VideoEncoder hwFallback = new VideoEncoder();

    private Callback callback;
    private final AtomicLong nativeHandle = new AtomicLong();
    private int width;
    private int height;
    private String mime;
    private byte[] yuvScratch;
    private boolean forceKeyframe;

    /** Active capture pipeline; {@link #releaseCapture()} swaps it out via {@code getAndSet(null)}. */
    private final AtomicReference<CapturePipeline> capturePipeline = new AtomicReference<>();

    public LibvpxVideoEncoderSink(CastSettings.VideoCodec preference) {
        this.preference = preference;
    }

    @Override
    public Surface getInputSurface() {
        CapturePipeline pipe = capturePipeline.get();
        if (nativeHandle.get() != 0L && pipe != null) {
            return pipe.reader.getSurface();
        }
        return hwFallback.getInputSurface();
    }

    @Override
    public void requestKeyframe() {
        if (nativeHandle.get() != 0L) {
            NativeCodecBridge.vpxRequestKeyframe(nativeHandle.get());
            forceKeyframe = true;
        } else {
            hwFallback.requestKeyframe();
        }
    }

    @Override
    public Surface prepare(int width, int height, String mime, CastTuningEngine.EffectiveParams params,
            Callback callback) throws IOException {
        releaseCapture();
        releaseNativeEncoder();
        this.callback = callback;
        this.width = width;
        this.height = height;
        this.mime = mime;
        yuvScratch = new byte[RgbToYuv420.bufferSize(width, height)];

        int bitrateKbps = Math.max(500, params.videoBitrate / 1000);
        long handle = NativeCodecBridge.vpxEncoderCreate(mime, width, height, bitrateKbps);
        if (handle == 0L) {
            Log.w(TAG, "libvpx encoder init failed; using MediaCodec for surface capture");
            CodecSessionRegistry.setVideoBackend(VideoCodecBackend.MEDIA_CODEC_HW, mime);
            return hwFallback.prepare(width, height, mime, params, callback);
        }
        nativeHandle.set(handle);

        CodecSessionRegistry.setVideoBackend(VideoCodecBackend.LIBVPX_NATIVE, mime);
        CapturePipeline pipe = CapturePipeline.create(width, height, this::onImageAvailable);
        capturePipeline.set(pipe);
        callback.onConfig(width, height, null, null);
        Log.i(TAG, "libvpx surface encoder " + width + "x" + height + " mime=" + mime);
        return pipe.reader.getSurface();
    }

    @Override
    public void prepareBufferInput(int width, int height, String mime, CastTuningEngine.EffectiveParams params,
            Callback callback) throws IOException {
        releaseCapture();
        releaseNativeEncoder();
        this.width = width;
        this.height = height;
        this.mime = mime;
        this.callback = callback;
        yuvScratch = new byte[RgbToYuv420.bufferSize(width, height)];
        int bitrateKbps = Math.max(500, params.videoBitrate / 1000);
        long handle = NativeCodecBridge.vpxEncoderCreate(mime, width, height, bitrateKbps);
        if (handle != 0L) {
            nativeHandle.set(handle);
            CodecSessionRegistry.setVideoBackend(VideoCodecBackend.LIBVPX_NATIVE, mime);
            callback.onConfig(width, height, null, null);
            return;
        }
        CodecSessionRegistry.setVideoBackend(VideoCodecBackend.MEDIA_CODEC_HW, mime);
        hwFallback.prepareBufferInput(width, height, mime, params, callback);
    }

    private void onImageAvailable(ImageReader reader) {
        CapturePipeline pipe = capturePipeline.get();
        if (pipe == null || pipe.reader != reader || callback == null) {
            return;
        }
        long handle = nativeHandle.get();
        if (handle == 0L) {
            return;
        }
        try (Image image = reader.acquireLatestImage()) {
            if (image == null) {
                return;
            }
            if (image.getWidth() != width || image.getHeight() != height) {
                return;
            }
            byte[] yuv = RgbToYuv420.fromRgbaImage(image, yuvScratch);
            if (yuv == null) {
                return;
            }
            if (capturePipeline.get() != pipe) {
                return;
            }
            long ptsUs = image.getTimestamp() / 1000L;
            emitEncodedFrame(handle, yuv, ptsUs);
        } catch (Exception e) {
            Log.w(TAG, "Capture frame skipped: " + e.getMessage());
        }
    }

    private void emitEncodedFrame(long handle, byte[] yuv, long ptsUs) {
        if (nativeHandle.get() != handle) {
            return;
        }
        VpxEncodedFrame frame = NativeCodecBridge.vpxEncodeYuvFrame(handle, yuv, ptsUs, forceKeyframe);
        forceKeyframe = false;
        if (frame != null && frame.data != null && callback != null) {
            callback.onEncodedFrame(ptsUs, frame.keyFrame, frame.data);
        }
    }

    @Override
    public void queueBitmapFrame(Bitmap bitmap, long ptsUs) {
        long handle = nativeHandle.get();
        if (handle == 0L) {
            hwFallback.queueBitmapFrame(bitmap, ptsUs);
            return;
        }
        byte[] yuv = RgbToYuv420.fromBitmap(bitmap, RgbToYuv420.Layout.I420);
        emitEncodedFrame(handle, yuv, ptsUs);
    }

    @Override
    public void queueYuvFrame(byte[] yuv, long ptsUs) {
        long handle = nativeHandle.get();
        if (handle == 0L) {
            hwFallback.queueYuvFrame(yuv, ptsUs);
            return;
        }
        emitEncodedFrame(handle, yuv, ptsUs);
    }

    @Override
    public void stop() {
        releaseCapture();
        releaseNativeEncoder();
        hwFallback.stop();
    }

    private void releaseCapture() {
        CapturePipeline old = capturePipeline.getAndSet(null);
        if (old != null) {
            old.release();
        }
    }

    private void releaseNativeEncoder() {
        long handle = nativeHandle.getAndSet(0L);
        if (handle != 0L) {
            NativeCodecBridge.vpxEncoderRelease(handle);
        }
    }

    private static final class CapturePipeline {
        final ImageReader reader;
        final HandlerThread thread;

        private CapturePipeline(ImageReader reader, HandlerThread thread) {
            this.reader = reader;
            this.thread = thread;
        }

        static CapturePipeline create(int width, int height, ImageReader.OnImageAvailableListener listener) {
            HandlerThread thread = new HandlerThread("libvpx-capture");
            thread.start();
            Handler handler = new Handler(thread.getLooper());
            ImageReader reader = ImageReader.newInstance(width, height, PixelFormat.RGBA_8888, 2);
            reader.setOnImageAvailableListener(listener, handler);
            return new CapturePipeline(reader, thread);
        }

        void release() {
            reader.setOnImageAvailableListener(null, null);
            reader.close();
            thread.quitSafely();
            try {
                thread.join(2000L);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }
    }

    static String mimeFor(CastSettings.VideoCodec preference) {
        if (preference == CastSettings.VideoCodec.LIBVPX_VP9) {
            return MediaFormat.MIMETYPE_VIDEO_VP9;
        }
        return MediaFormat.MIMETYPE_VIDEO_VP8;
    }
}
