package com.foxx.androidcast.sender;

/*********************************************************************
 * VideoEncoder.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 VideoEncoder
 * Contributors:
 *    - Anton Afanasyeu <a.afanasieff@gmail.com> (8 commits, 279 lines)
 *    - Cursor Agent (project assistant)
 * Digest: SHA256	40dcb27fe239420b524f86cc5cd367356bfa116b23b59ce08ef4d81a71d3dcbb
**********************************************************************/
import android.graphics.Bitmap;
import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import android.os.Build;
import android.os.Bundle;
import android.util.Log;
import android.view.Surface;

import com.foxx.androidcast.CastTuningEngine;
import com.foxx.androidcast.media.RgbToYuv420;
import com.foxx.androidcast.media.codec.VideoEncoderSink;

import java.io.IOException;
import java.nio.ByteBuffer;

/** Hardware video encoder (H.264 / H.265 / VP8 / VP9) via MediaCodec. */
public class VideoEncoder implements VideoEncoderSink {
    private static final String TAG = "VideoEncoder";

    private MediaCodec codec;
    private Surface inputSurface;
    private boolean bufferInput;
    private RgbToYuv420.Layout yuvLayout = RgbToYuv420.Layout.I420;
    private byte[] yuvScratch;
    private Callback callback;
    private int width;
    private int height;
    private volatile boolean running;
    private Thread drainThread;
    private String mime;

    public Surface getInputSurface() {
        return inputSurface;
    }

    public void requestKeyframe() {
        if (codec == null || !running) {
            return;
        }
        try {
            Bundle params = new Bundle();
            params.putInt(MediaCodec.PARAMETER_KEY_REQUEST_SYNC_FRAME, 0);
            codec.setParameters(params);
        } catch (Exception e) {
            Log.w(TAG, "Keyframe request failed: " + e.getMessage());
        }
    }

    public Surface prepare(int width, int height, String mime, CastTuningEngine.EffectiveParams params,
            Callback callback) throws IOException {
        return prepareInternal(width, height, mime, params, callback, false);
    }

    /** Software-fed frames (calibration chart) via YUV buffers — no input Surface. */
    public void prepareBufferInput(int width, int height, String mime, CastTuningEngine.EffectiveParams params,
            Callback callback) throws IOException {
        prepareInternal(width, height, mime, params, callback, true);
    }

    private Surface prepareInternal(int width, int height, String mime, CastTuningEngine.EffectiveParams params,
            Callback callback, boolean buffer) throws IOException {
        this.width = width;
        this.height = height;
        this.mime = mime;
        this.callback = callback;
        this.bufferInput = buffer;
        yuvScratch = buffer ? new byte[RgbToYuv420.bufferSize(width, height)] : null;

        codec = MediaCodec.createEncoderByType(mime);
        yuvLayout = buffer ? pickYuvLayout(codec, mime) : RgbToYuv420.Layout.I420;

        MediaFormat tuned = buildFormat(width, height, mime, params, true, buffer, yuvLayout);
        MediaFormat basic = buildFormat(width, height, mime, params, false, buffer, yuvLayout);

        if (!tryConfigure(codec, tuned)) {
            Log.w(TAG, "Tuned encoder format rejected; retrying with basic format");
            codec.release();
            codec = MediaCodec.createEncoderByType(mime);
            yuvLayout = buffer ? pickYuvLayout(codec, mime) : RgbToYuv420.Layout.I420;
            basic = buildFormat(width, height, mime, params, false, buffer, yuvLayout);
            if (!tryConfigure(codec, basic)) {
                codec.release();
                codec = null;
                throw new IOException("Encoder configure failed for " + mime + " " + width + "x" + height);
            }
        }

        if (bufferInput) {
            inputSurface = null;
        } else {
            inputSurface = codec.createInputSurface();
        }
        codec.start();
        running = true;
        drainThread = new Thread(this::drainLoop, "VideoEncoderDrain");
        drainThread.start();
        Log.i(TAG, (bufferInput ? "Buffer/" + yuvLayout : "Surface") + " encoder " + mime
                + " " + width + "x" + height + " @ " + params.videoFps + "fps");
        return inputSurface;
    }

    private static RgbToYuv420.Layout pickYuvLayout(MediaCodec codec, String mime) {
        MediaCodecInfo.CodecCapabilities caps = codec.getCodecInfo().getCapabilitiesForType(mime);
        for (int fmt : caps.colorFormats) {
            if (fmt == MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar) {
                return RgbToYuv420.Layout.NV12;
            }
        }
        return RgbToYuv420.Layout.I420;
    }

    public void queueBitmapFrame(Bitmap bitmap, long ptsUs) {
        if (!bufferInput || codec == null || !running || bitmap == null) {
            return;
        }
        byte[] yuv = RgbToYuv420.fromBitmap(bitmap, yuvLayout);
        queueYuvFrame(yuv, ptsUs);
    }

    public void queueYuvFrame(byte[] yuv, long ptsUs) {
        if (!bufferInput || codec == null || !running || yuv == null) {
            return;
        }
        try {
            int index = codec.dequeueInputBuffer(0);
            if (index < 0) {
                return;
            }
            ByteBuffer buffer = codec.getInputBuffer(index);
            if (buffer == null) {
                return;
            }
            buffer.clear();
            int len = Math.min(buffer.remaining(), yuv.length);
            buffer.put(yuv, 0, len);
            codec.queueInputBuffer(index, 0, len, ptsUs, 0);
        } catch (Exception e) {
            Log.w(TAG, "queueYuvFrame: " + e.getMessage());
        }
    }

    private static MediaFormat buildFormat(int width, int height, String mime,
            CastTuningEngine.EffectiveParams params, boolean tuned, boolean buffer,
            RgbToYuv420.Layout layout) {
        MediaFormat format = MediaFormat.createVideoFormat(mime, width, height);
        int color;
        if (!buffer) {
            color = MediaCodecInfo.CodecCapabilities.COLOR_FormatSurface;
        } else if (layout == RgbToYuv420.Layout.NV12) {
            color = MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar;
        } else {
            color = MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Planar;
        }
        format.setInteger(MediaFormat.KEY_COLOR_FORMAT, color);
        format.setInteger(MediaFormat.KEY_BIT_RATE, params.videoBitrate);
        format.setInteger(MediaFormat.KEY_FRAME_RATE, params.videoFps);
        format.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, params.keyframeIntervalSec);
        if (!tuned || !params.requestLowLatency) {
            return format;
        }
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.LOLLIPOP) {
            format.setInteger(MediaFormat.KEY_BITRATE_MODE,
                    MediaCodecInfo.EncoderCapabilities.BITRATE_MODE_CBR);
        }
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.M) {
            format.setInteger(MediaFormat.KEY_PRIORITY, 0);
        }
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.N) {
            format.setInteger(MediaFormat.KEY_OPERATING_RATE, params.videoFps);
        }
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R) {
            format.setInteger(MediaFormat.KEY_LOW_LATENCY, 1);
        }
        return format;
    }

    private static boolean tryConfigure(MediaCodec codec, MediaFormat format) {
        try {
            codec.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
            Log.i(TAG, "Using encoder format: " + format);
            return true;
        } catch (IllegalArgumentException | IllegalStateException e) {
            Log.w(TAG, "Encoder configure rejected: " + format + " — " + e.getMessage());
            return false;
        }
    }

    public void stop() {
        running = false;
        bufferInput = false;
        if (drainThread != null) {
            try {
                drainThread.join(1500);
            } catch (InterruptedException ignored) {
                Thread.currentThread().interrupt();
            }
        }
        if (codec != null) {
            try {
                codec.stop();
            } catch (Exception ignored) {
            }
            codec.release();
            codec = null;
        }
        if (inputSurface != null) {
            inputSurface.release();
            inputSurface = null;
        }
        yuvScratch = null;
    }

    private void drainLoop() {
        MediaCodec.BufferInfo info = new MediaCodec.BufferInfo();
        boolean configSent = false;
        while (running) {
            MediaCodec c = codec;
            if (c == null) {
                break;
            }
            int index = c.dequeueOutputBuffer(info, 5_000);
            if (index == MediaCodec.INFO_TRY_AGAIN_LATER) {
                continue;
            }
            if (index == MediaCodec.INFO_OUTPUT_FORMAT_CHANGED) {
                sendConfig(c.getOutputFormat(), configSent);
                configSent = true;
                continue;
            }
            if (index < 0) {
                continue;
            }
            ByteBuffer buffer = c.getOutputBuffer(index);
            if (buffer != null && info.size > 0) {
                byte[] data = new byte[info.size];
                buffer.position(info.offset);
                buffer.limit(info.offset + info.size);
                buffer.get(data);
                boolean keyFrame = (info.flags & MediaCodec.BUFFER_FLAG_CODEC_CONFIG) == 0
                        && (info.flags & MediaCodec.BUFFER_FLAG_KEY_FRAME) != 0;
                if ((info.flags & MediaCodec.BUFFER_FLAG_CODEC_CONFIG) == 0 && callback != null) {
                    callback.onEncodedFrame(info.presentationTimeUs, keyFrame, data);
                }
            }
            c.releaseOutputBuffer(index, false);
            if ((info.flags & MediaCodec.BUFFER_FLAG_END_OF_STREAM) != 0) {
                break;
            }
        }
    }

    private void sendConfig(MediaFormat format, boolean alreadySent) {
        if (alreadySent || callback == null) {
            return;
        }
        ByteBuffer sps = format.getByteBuffer("csd-0");
        ByteBuffer pps = format.getByteBuffer("csd-1");
        byte[] csd0 = sps != null ? toArray(sps) : new byte[0];
        byte[] csd1 = pps != null ? toArray(pps) : new byte[0];
        callback.onConfig(width, height, csd0, csd1);
    }

    private static byte[] toArray(ByteBuffer buffer) {
        byte[] arr = new byte[buffer.remaining()];
        ByteBuffer dup = buffer.duplicate();
        dup.get(arr);
        return arr;
    }
}
