package com.ttstd.controlled.webrtc;

import android.media.Image;
import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaCodecList;
import android.media.MediaFormat;
import android.os.Bundle;
import android.util.Log;

import org.webrtc.DataChannel;
import org.webrtc.VideoFrame;

import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.ArrayDeque;

/**
 * 自建屏幕编码通道：复用 WebRTC 的屏幕采集管线（ScreenCapturerAndroid 的 VideoFrame），
 * 把每一帧编码为 H.264 裸流后，经 WebRTC 的 video DataChannel 分片透传给控制端
 * （控制端使用自建 MediaCodec 解码渲染）。
 *
 * 设计要点（Android 14 兼容）：
 * 在 Android 14+ 上，MediaProjection 授权令牌为一次性，且一个 MediaProjection 实例
 * 只能 createVirtualDisplay 一次。因此自编码模式【不再】自己申请 MediaProjection /
 * VirtualDisplay，而是复用 WebRTC 已经建立的、唯一的屏幕采集管线，从 VideoFrame 直接编码。
 * 这样全程只有一个虚拟显示、一个令牌，既避免了
 * "Must register a callback before starting capture"，
 * 也避免了 "Don't re-use the resultData ... Don't take multiple captures ..." 两类崩溃，
 * 同时运行时可在 WebRTC 与自编码模式间自由切换（共享同一条采集）。
 */
public class SelfCodecEncoder {

    public static final int STREAM_MODE_WEBRTC = 0;     // WebRTC 内置媒体流
    public static final int STREAM_MODE_SELF_CODEC = 1;  // 自编码（本类）

    // ---- 自定义二进制协议（与控制端 SelfCodecDecoder 保持一致）----
    private static final byte MAGIC = (byte) 0xAB;
    private static final int CHUNK_SIZE = 16384;         // 每个 DataChannel 分片大小
    private static final byte UNIT_TYPE_CONFIG = 1;       // SPS/PPS 等参数集
    private static final byte UNIT_TYPE_FRAME = 2;        // 编码帧

    private static final String TAG = "SelfCodecEncoder";

    private final int width;
    private final int height;
    private final int fps;
    private final DataChannel videoChannel;
    private final Callback callback;

    private MediaCodec encoder;
    private boolean running = false;
    private int frameSeq = 0;
    private final Object sendLock = new Object();
    private final Object inputLock = new Object();
    private final ArrayDeque<Integer> freeInputIndices = new ArrayDeque<>();
    private byte[] nv12Buffer;
    // 缓存 SPS/PPS 等参数集：video 通道不可靠（可能丢包），在每个关键帧前重发，
    // 保证控制端即使错过首次 CONFIG 也能在下一个 IDR 前完成解码器配置。
    private final java.util.ArrayList<byte[]> cachedConfigs = new java.util.ArrayList<>();

    public interface Callback {
        void onEncoderStarted(int w, int h, int fps);

        void onEncoderError(String error);
    }

    public SelfCodecEncoder(int width, int height, int fps, DataChannel videoChannel, Callback cb) {
        this.width = width;
        this.height = height;
        this.fps = fps;
        this.videoChannel = videoChannel;
        this.callback = cb;
    }

    public void start() {
        if (running) return;
        try {
            MediaFormat format = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC, width, height);
            int bitrate = computeBitrate(width, height, fps);
            format.setInteger(MediaFormat.KEY_BIT_RATE, bitrate);
            format.setInteger(MediaFormat.KEY_FRAME_RATE, fps);
            format.setInteger(MediaFormat.KEY_CAPTURE_RATE, fps);
            format.setInteger(MediaFormat.KEY_COLOR_FORMAT, selectColorFormat());
            // 关键帧间隔：1 秒一个 IDR，便于断流后快速恢复；也可经 requestKeyFrame 即时强求
            format.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 1);

            encoder = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_AVC);
            // 异步回调需在 configure() 之前设置（部分系统版本在 configure 后设置会抛异常）
            encoder.setCallback(encoderCallback);
            encoder.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
            encoder.start();

            nv12Buffer = new byte[width * height * 3 / 2];
            running = true;
            if (callback != null) callback.onEncoderStarted(width, height, fps);
            Log.i(TAG, "self codec encoder started: " + width + "x" + height + "@" + fps + ", bitrate=" + bitrate);
        } catch (Exception e) {
            Log.e(TAG, "start encoder failed", e);
            release();
            if (callback != null) callback.onEncoderError(e.getMessage());
        }
    }

    /** 选择合适的 YUV 输入颜色格式（优先 NV12 系；几乎全部 AVC 编码器均支持）。 */
    private int selectColorFormat() {
        try {
            MediaCodecList list = new MediaCodecList(MediaCodecList.REGULAR_CODECS);
            MediaCodecInfo info = null;
            for (MediaCodecInfo c : list.getCodecInfos()) {
                if (!c.isEncoder()) continue;
                for (String t : c.getSupportedTypes()) {
                    if (t.equalsIgnoreCase(MediaFormat.MIMETYPE_VIDEO_AVC)) {
                        info = c;
                        break;
                    }
                }
                if (info != null) break;
            }
            if (info == null) return MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar;
            MediaCodecInfo.CodecCapabilities caps = info.getCapabilitiesForType(MediaFormat.MIMETYPE_VIDEO_AVC);
            // 优先 NV12（COLOR_FormatYUV420SemiPlanar）：其输入 Image 为 2 平面（Y + UV 交错），
            // 与下方 writeNv12 的写入逻辑严格一致。仅在设备不支持 NV12 时才退而求其次用 Flexible。
            for (int f : caps.colorFormats) {
                if (f == MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar) {
                    return f;
                }
            }
            for (int f : caps.colorFormats) {
                if (f == MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Flexible) {
                    return f;
                }
            }
            return MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar;
        } catch (Exception e) {
            return MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar;
        }
    }

    private int computeBitrate(int w, int h, int fps) {
        int b = (int) (w * h * fps * 0.12);
        return Math.max(500_000, Math.min(8_000_000, b));
    }

    /**
     * 由 WebRTC 采集管线喂入一帧（共享同一条屏幕采集，无需再申请 VirtualDisplay）。
     * 当编码器输入缓冲不足时会丢弃该帧（编码器跟不上），保证不堆积、不阻塞采集线程。
     */
    public void feedFrame(VideoFrame frame) {
        synchronized (this) {
            MediaCodec mc = encoder;
            if (!running || mc == null) return;

            // 检查分辨率：自编码器 MediaCodec 尺寸是固定的，若采集分辨率已切换则丢弃本帧，
            // 等待外部（ScreenCaptureService）重启本编码器实例。
            if (frame.getRotatedWidth() != width || frame.getRotatedHeight() != height) {
                return;
            }

            Integer idx;
            synchronized (inputLock) {
                idx = freeInputIndices.poll();
            }
            if (idx == null) {
                return; // 丢掉这一帧，等待下一个空闲输入缓冲
            }
            try {
                encodeFrame(mc, idx, frame);
            } catch (Exception e) {
                Log.e(TAG, "encodeFrame failed", e);
                try {
                    mc.queueInputBuffer(idx, 0, 0, System.nanoTime() / 1000, 0);
                } catch (Exception ignored) {
                }
            }
        }
    }

    private void encodeFrame(MediaCodec mc, int index, VideoFrame frame) {
        long ptsUs = System.nanoTime() / 1000;
        VideoFrame.I420Buffer i420 = frame.getBuffer().toI420();
        if (i420 == null) {
            mc.queueInputBuffer(index, 0, 0, ptsUs, 0);
            return;
        }
        int w = i420.getWidth();
        int h = i420.getHeight();
        byte[] buffer = ensureNv12Buffer(w, h);
        i420ToNv12(i420, buffer, w, h);
        i420.release();
        writeNv12(mc, index, buffer, w, h, ptsUs);
    }

    private byte[] ensureNv12Buffer(int w, int h) {
        int need = w * h * 3 / 2;
        byte[] buf = nv12Buffer;
        if (buf == null || buf.length < need) {
            buf = new byte[need];
            nv12Buffer = buf;
        }
        return buf;
    }

    /** I420（三平面）→ NV12（Y + UV 交错），写入 out。 */
    private void i420ToNv12(VideoFrame.I420Buffer i420, byte[] out, int w, int h) {
        ByteBuffer y = i420.getDataY();
        ByteBuffer u = i420.getDataU();
        ByteBuffer v = i420.getDataV();
        int yStride = i420.getStrideY();
        int uStride = i420.getStrideU();
        int vStride = i420.getStrideV();
        int off = 0;
        for (int r = 0; r < h; r++) {
            y.position(r * yStride);
            y.get(out, off, w);
            off += w;
        }
        int cw = w / 2;
        int ch = h / 2;
        for (int r = 0; r < ch; r++) {
            for (int c = 0; c < cw; c++) {
                out[off++] = u.get(r * uStride + c);
                out[off++] = v.get(r * vStride + c);
            }
        }
    }

    /** 把 NV12 按编码器输入 Image 的 stride 写入输入缓冲并入队。 */
    private void writeNv12(MediaCodec mc, int index, byte[] nv12, int w, int h, long ptsUs) {
        // 使用 try-with-resources 确保 Image 及时关闭。
        // 根据 MediaCodec 文档，Image 必须在 queueInputBuffer 之前关闭，否则可能导致 native 崩溃或阻塞。
        try (Image image = mc.getInputImage(index)) {
            if (image == null) {
                mc.queueInputBuffer(index, 0, 0, ptsUs, 0);
                return;
            }
            Image.Plane[] planes = image.getPlanes();
            // Y 平面写入
            Image.Plane yPlane = planes[0];
            ByteBuffer yBuf = yPlane.getBuffer();
            int yRow = yPlane.getRowStride();
            int yPix = yPlane.getPixelStride();
            int yLimit = yBuf.limit();

            int src = 0;
            for (int r = 0; r < h; r++) {
                int dstRow = r * yRow;
                for (int c = 0; c < w; c++) {
                    int dst = dstRow + c * yPix;
                    if (dst < yLimit) {
                        yBuf.put(dst, nv12[src + c]);
                    }
                }
                src += w;
            }

            // UV 平面写入
            int ch = h / 2;
            int cw = w / 2;
            int uvSrcBase = w * h;

            if (planes.length >= 3 && planes[1].getPixelStride() == 1 && planes[2].getPixelStride() == 1) {
                // I420 平面格式 (Planar)
                for (int p = 1; p <= 2; p++) {
                    Image.Plane plane = planes[p];
                    ByteBuffer buf = plane.getBuffer();
                    int row = plane.getRowStride();
                    int limit = buf.limit();
                    for (int r = 0; r < ch; r++) {
                        int dstRow = r * row;
                        for (int c = 0; c < cw; c++) {
                            int s = uvSrcBase + (r * cw + c) * 2 + (p - 1);
                            int dst = dstRow + c;
                            if (dst < limit) {
                                buf.put(dst, nv12[s]);
                            }
                        }
                    }
                }
            } else if (planes.length >= 2) {
                // NV12/NV21 半平面格式 (Semi-planar)
                Image.Plane uvPlane = planes[1];
                ByteBuffer uvBuf = uvPlane.getBuffer();
                int uvRow = uvPlane.getRowStride();
                int uvPix = uvPlane.getPixelStride();
                int uvLimit = uvBuf.limit();

                for (int r = 0; r < ch; r++) {
                    int dRow = r * uvRow;
                    int sRow = uvSrcBase + r * cw * 2;
                    for (int c = 0; c < cw; c++) {
                        int s = sRow + c * 2;
                        int d = dRow + c * uvPix;
                        if (d < uvLimit) {
                            uvBuf.put(d, nv12[s]);
                        }
                        if (d + 1 < uvLimit) {
                            uvBuf.put(d + 1, nv12[s + 1]);
                        }
                    }
                }
            }
            // 显式关闭 image 确保在 queueInputBuffer 前释放 native 句柄
            image.close();
            mc.queueInputBuffer(index, 0, w * h * 3 / 2, ptsUs, 0);
        } catch (Exception e) {
            Log.e(TAG, "writeNv12 failed", e);
            try {
                mc.queueInputBuffer(index, 0, 0, ptsUs, 0);
            } catch (Exception ignored) {
            }
        }
    }

    private final MediaCodec.Callback encoderCallback = new MediaCodec.Callback() {
        @Override
        public void onInputBufferAvailable(MediaCodec mc, int index) {
            synchronized (inputLock) {
                freeInputIndices.add(index);
            }
        }

        @Override
        public void onOutputFormatChanged(MediaCodec mc, MediaFormat format) {
            // 主流设备在此给出 SPS/PPS，作为 CONFIG 单元发送
            sendCsds(format);
        }

        @Override
        public void onOutputBufferAvailable(MediaCodec mc, int index, MediaCodec.BufferInfo info) {
            try {
                ByteBuffer buf = mc.getOutputBuffer(index);
                if (buf == null) {
                    mc.releaseOutputBuffer(index, false);
                    return;
                }
                byte[] data = new byte[info.size];
                buf.position(info.offset);
                buf.get(data);
                boolean isKey = (info.flags & MediaCodec.BUFFER_FLAG_KEY_FRAME) != 0;
                boolean isConfig = (info.flags & MediaCodec.BUFFER_FLAG_CODEC_CONFIG) != 0;
                long pts = info.presentationTimeUs;
                mc.releaseOutputBuffer(index, false);

                if (isConfig) {
                    cacheConfig(data);
                    sendUnit(UNIT_TYPE_CONFIG, 0, false, data);
                } else {
                    if (isKey) {
                        // 关键帧前重发参数集：解码端已配置时会忽略重复 CONFIG，
                        // 未配置（首次 CONFIG 丢失/通道未就绪）时借此完成配置。
                        resendCachedConfigs();
                    }
                    sendUnit(UNIT_TYPE_FRAME, pts, isKey, data);
                }
            } catch (Exception e) {
                Log.e(TAG, "onOutputBufferAvailable error", e);
            }
        }

        @Override
        public void onError(MediaCodec mc, MediaCodec.CodecException e) {
            Log.e(TAG, "encoder error", e);
            if (callback != null) callback.onEncoderError(e.getMessage());
        }
    };

    private void sendCsds(MediaFormat format) {
        if (format == null) return;
        ByteBuffer sps = format.getByteBuffer("csd-0");
        ByteBuffer pps = format.getByteBuffer("csd-1");
        if (sps != null) {
            byte[] d = toBytes(sps);
            cacheConfig(d);
            sendUnit(UNIT_TYPE_CONFIG, 0, false, d);
        }
        if (pps != null) {
            byte[] d = toBytes(pps);
            cacheConfig(d);
            sendUnit(UNIT_TYPE_CONFIG, 0, false, d);
        }
    }

    /** 缓存参数集（去重），供关键帧前重发。 */
    private void cacheConfig(byte[] data) {
        if (data == null || data.length == 0) return;
        synchronized (cachedConfigs) {
            for (byte[] c : cachedConfigs) {
                if (java.util.Arrays.equals(c, data)) return;
            }
            cachedConfigs.add(data);
        }
    }

    /** 在每个关键帧前重发缓存的参数集。 */
    private void resendCachedConfigs() {
        byte[][] snapshot;
        synchronized (cachedConfigs) {
            if (cachedConfigs.isEmpty()) return;
            snapshot = cachedConfigs.toArray(new byte[0][]);
        }
        for (byte[] c : snapshot) {
            sendUnit(UNIT_TYPE_CONFIG, 0, false, c);
        }
    }

    private static byte[] toBytes(ByteBuffer b) {
        ByteBuffer d = b.duplicate();
        byte[] out = new byte[d.remaining()];
        d.get(out);
        return out;
    }

    /** 组装一个"单元"（类型 + pts + isKey + 数据），再切片经 videoChannel 发送。 */
    private void sendUnit(byte unitType, long pts, boolean isKey, byte[] data) {
        ByteBuffer unit = ByteBuffer.allocate(1 + 4 + 1 + 4 + data.length).order(ByteOrder.BIG_ENDIAN);
        unit.put(unitType);
        unit.putInt((int) (pts & 0xFFFFFFFFL));
        unit.put((byte) (isKey ? 1 : 0));
        unit.putInt(data.length);
        unit.put(data);
        byte[] unitBytes = unit.array();

        int totalChunks = (unitBytes.length + CHUNK_SIZE - 1) / CHUNK_SIZE;
        if (totalChunks == 0) totalChunks = 1;
        int seq;
        synchronized (sendLock) {
            seq = frameSeq++;
        }
        for (int i = 0; i < totalChunks; i++) {
            int off = i * CHUNK_SIZE;
            int len = Math.min(CHUNK_SIZE, unitBytes.length - off);
            byte[] chunk = new byte[len];
            System.arraycopy(unitBytes, off, chunk, 0, len);
            sendChunk(seq, (short) totalChunks, (short) i, chunk);
        }
    }

    private void sendChunk(int seq, short total, short idx, byte[] chunk) {
        if (videoChannel == null || videoChannel.state() != DataChannel.State.OPEN) return;
        ByteBuffer msg = ByteBuffer.allocate(1 + 4 + 2 + 2 + 4 + chunk.length).order(ByteOrder.BIG_ENDIAN);
        msg.put(MAGIC);
        msg.putInt(seq);
        msg.putShort(total);
        msg.putShort(idx);
        msg.putInt(chunk.length);
        msg.put(chunk);
        msg.flip();
        try {
            videoChannel.send(new DataChannel.Buffer(msg, true));
        } catch (Exception e) {
            Log.e(TAG, "send video chunk failed", e);
        }
    }

    /** 即时请求一个关键帧（强求 IDR）。 */
    public void requestKeyFrame() {
        MediaCodec mc = encoder;
        if (mc == null) return;
        try {
            Bundle params = new Bundle();
            params.putInt(MediaCodec.PARAMETER_KEY_REQUEST_SYNC_FRAME, 0);
            mc.setParameters(params);
        } catch (Exception e) {
            Log.w(TAG, "requestKeyFrame failed", e);
        }
    }

    /** 返回编码器是否正在运行（供上层判断启动结果）。 */
    public boolean isRunning() {
        return running;
    }

    public void release() {
        synchronized (this) {
            running = false;
            freeInputIndices.clear();
            if (encoder != null) {
                try {
                    encoder.stop();
                } catch (Exception ignored) {
                }
                try {
                    encoder.release();
                } catch (Exception ignored) {
                }
                encoder = null;
            }
            nv12Buffer = null;
        }
    }
}
