package com.ttstd.controller.webrtc;

import android.media.MediaCodec;
import android.media.MediaFormat;
import android.os.Build;
import android.util.Log;
import android.view.Surface;

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

/**
 * 自建屏幕解码通道：从 WebRTC 的 video DataChannel 收取被控端发来的 H.264 裸流分片，
 * 重组后使用 MediaCodec 解码并渲染到指定 Surface。与 WebRTC 默认媒体轨道解耦，
 * 配合被控端 SelfCodecEncoder 使用。
 */
public class SelfCodecDecoder {

    private static final String TAG = "SelfCodecDecoder";

    // ---- 必须与被控端 SelfCodecEncoder 保持一致 ----
    private static final byte MAGIC = (byte) 0xAB;
    private static final int CHUNK_SIZE = 16384;
    private static final int UNIT_TYPE_CONFIG = 1;
    private static final int UNIT_TYPE_FRAME = 2;
    private static final int PLACEHOLDER_W = 1920;
    private static final int PLACEHOLDER_H = 1080;

    private MediaCodec decoder;
    private Surface surface;
    private boolean configured = false;
    private boolean enabled = false;

    // 实际宽高（由被控端 REPORT_RESOLUTION 上报），若未上报则使用 PLACEHOLDER
    private int videoWidth = PLACEHOLDER_W;
    private int videoHeight = PLACEHOLDER_H;

    // 分辨率变化回调（通知 UI 调整渲染 Surface 尺寸以保持宽高比，避免画面被拉伸变形）
    private OnResolutionUpdateListener resolutionListener;

    // 待喂入的参数集（SPS/PPS 等），以 CODEC_CONFIG 形式输入解码器
    private final ArrayDeque<byte[]> pendingConfigs = new ArrayDeque<>();
    // 待解码帧
    private final ArrayDeque<Frame> frameQueue = new ArrayDeque<>();
    // 解码器可用的输入缓冲下标
    private final ArrayDeque<Integer> inputQueue = new ArrayDeque<>();
    // 分片重组：frameSeq -> 该单元的若干分片
    private final Map<Integer, ChunkBuffer> assembling = new HashMap<>();
    private int lastSeq = -1;

    private static class Frame {
        byte[] data;
        long pts;
        int flags;
    }

    private static class ChunkBuffer {
        int total;
        byte[][] chunks;
        int received;
    }

    public void setSurface(Surface s) {
        this.surface = s;
        tryConfigure();
    }

    public void setEnabled(boolean e) {
        enabled = e;
        if (!e) {
            releaseDecoder();
        }
    }

    /**
     * 更新当前视频流的实际宽高（由外部 WebRtcClient 收到被控端上报后调用）。
     * 如果宽高发生变化且解码器已配置，则重置解码器以应用新分辨率，避免黑屏或花屏。
     */
    public void updateResolution(int w, int h) {
        if (w <= 0 || h <= 0) return;
        if (this.videoWidth != w || this.videoHeight != h) {
            Log.i(TAG, "Resolution update: " + videoWidth + "x" + videoHeight + " -> " + w + "x" + h);
            this.videoWidth = w;
            this.videoHeight = h;
            if (configured) {
                // 分辨率变化时需重置解码器以重新 configure
                releaseDecoder();
            }
            if (resolutionListener != null) {
                resolutionListener.onResolutionUpdated(w, h);
            }
        }
    }

    /** 收到 video DataChannel 的二进制消息（一条分片）。 */
    public void onBinaryMessage(ByteBuffer raw) {
        if (!enabled) return;
        if (raw == null) return;
        ByteBuffer b = raw.duplicate();
        if (b.remaining() < 13) return;
        byte magic = b.get();
        if (magic != MAGIC) return;
        int seq = b.getInt();
        // 过期检查：被控端每个“单元”共用同一个 seq（同一单元的所有分片 seq 相同、以 idx 区分），
        // 因此不能用 seq<=lastSeq 来丢弃——否则同一单元除首片外的分片都会被误删，
        // 导致大帧永远无法重组、自编码串流黑屏。这里只在 seq 明显小于 lastSeq（旧单元）时丢弃，
        // 相同 seq（同单元的其它分片）放行。
        if (lastSeq != -1 && seq < lastSeq && (lastSeq - seq) < 1000) {
            return;
        }
        short total = b.getShort();
        short idx = b.getShort();
        int len = b.getInt();
        if (b.remaining() < len) return;
        byte[] chunk = new byte[len];
        b.get(chunk);

        synchronized (assembling) {
            ChunkBuffer cb = assembling.get(seq);
            if (cb == null) {
                cb = new ChunkBuffer();
                cb.total = total;
                cb.chunks = new byte[total][];
                cb.received = 0;
                assembling.put(seq, cb);
                if (assembling.size() > 24) {
                    int min = Integer.MAX_VALUE;
                    for (int k : assembling.keySet()) {
                        if (k < min) min = k;
                    }
                    assembling.remove(min);
                }
            }
            if (idx >= 0 && idx < total && cb.chunks[idx] == null) {
                cb.chunks[idx] = chunk;
                cb.received++;
            }
            if (cb.received == total) {
                assembling.remove(seq);
                if (seq > lastSeq) lastSeq = seq;
                byte[] unit = join(cb);
                if (unit != null) handleUnit(unit);
            }
        }
    }

    private byte[] join(ChunkBuffer cb) {
        int totalLen = 0;
        for (byte[] c : cb.chunks) {
            if (c != null) totalLen += c.length;
        }
        ByteBuffer out = ByteBuffer.allocate(totalLen);
        for (byte[] c : cb.chunks) {
            if (c != null) out.put(c);
        }
        return out.array();
    }

    private void handleUnit(byte[] unit) {
        ByteBuffer b = ByteBuffer.wrap(unit);
        byte type = b.get();
        long pts = b.getInt() & 0xFFFFFFFFL;
        byte isKey = b.get();
        int len = b.getInt();
        byte[] data = new byte[len];
        b.get(data);

        if (type == UNIT_TYPE_CONFIG) {
            synchronized (pendingConfigs) {
                // 仅在尚未包含该配置时添加，避免关键帧前的重复重发堆积。
                // 关键：不因为 configured 而 return，确保 SPS 和 PPS 都能被加入队列。
                boolean exists = false;
                for (byte[] c : pendingConfigs) {
                    if (java.util.Arrays.equals(c, data)) {
                        exists = true;
                        break;
                    }
                }
                if (!exists) {
                    pendingConfigs.add(data);
                }
            }
            if (!configured) {
                tryConfigure();
            } else {
                feed(); // 已配置时也尝试喂入，以应对中途可能的参数集更新
            }
        } else if (type == UNIT_TYPE_FRAME) {
            Frame f = new Frame();
            f.data = data;
            f.pts = pts;
            f.flags = (isKey != 0) ? MediaCodec.BUFFER_FLAG_KEY_FRAME : 0;
            synchronized (frameQueue) {
                frameQueue.add(f);
            }
            feed();
        }
    }

    private void tryConfigure() {
        if (configured) return;
        if (surface == null) return;
        synchronized (pendingConfigs) {
            if (pendingConfigs.isEmpty()) return;
        }
        try {
            MediaFormat format = MediaFormat.createVideoFormat(
                    MediaFormat.MIMETYPE_VIDEO_AVC, videoWidth, videoHeight);
            // 关键：对部分芯片组（如 MTK/Exynos），低延迟标志有助于减少首帧黑屏时间（API 30+）
            if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R) {
                format.setInteger(MediaFormat.KEY_LOW_LATENCY, 1);
            }
            decoder = MediaCodec.createDecoderByType(MediaFormat.MIMETYPE_VIDEO_AVC);
            // 异步回调必须在 configure() 之前设置，否则抛 IllegalStateException，
            // 解码器将永远无法进入工作状态（黑屏）。
            decoder.setCallback(decoderCallback);
            decoder.configure(format, surface, null, 0);
            decoder.start();
            configured = true;
            Log.i(TAG, "self codec decoder configured");
            feed();
        } catch (Exception e) {
            Log.e(TAG, "decoder configure failed", e);
            // 清理半初始化的解码器，等待下一个 CONFIG 到达后重试
            if (decoder != null) {
                try {
                    decoder.release();
                } catch (Exception ignored) {
                }
                decoder = null;
            }
            configured = false;
        }
    }

    private final MediaCodec.Callback decoderCallback = new MediaCodec.Callback() {
        @Override
        public void onInputBufferAvailable(MediaCodec mc, int index) {
            synchronized (inputQueue) {
                inputQueue.add(index);
            }
            feed();
        }

        @Override
        public void onOutputBufferAvailable(MediaCodec mc, int index, MediaCodec.BufferInfo info) {
            try {
                mc.releaseOutputBuffer(index, true);
            } catch (IllegalStateException e) {
                Log.w(TAG, "onOutputBufferAvailable: codec might be stopped/released", e);
            }
        }

        @Override
        public void onError(MediaCodec mc, MediaCodec.CodecException e) {
            Log.e(TAG, "decoder error", e);
        }

        @Override
        public void onOutputFormatChanged(MediaCodec mc, MediaFormat format) {
            Log.d(TAG, "decoder output format changed");
        }
    };

    private void feed() {
        if (decoder == null || !configured) return;
        synchronized (pendingConfigs) {
            synchronized (inputQueue) {
                while (!pendingConfigs.isEmpty() && !inputQueue.isEmpty()) {
                    byte[] cfg = pendingConfigs.poll();
                    int idx = inputQueue.poll();
                    try {
                        ByteBuffer ib = decoder.getInputBuffer(idx);
                        ib.clear();
                        ib.put(cfg);
                        decoder.queueInputBuffer(idx, 0, cfg.length, 0, MediaCodec.BUFFER_FLAG_CODEC_CONFIG);
                    } catch (Exception e) {
                        Log.e(TAG, "queue config failed", e);
                        pendingConfigs.addFirst(cfg);
                        inputQueue.addFirst(idx);
                        break;
                    }
                }
            }
        }
        synchronized (frameQueue) {
            synchronized (inputQueue) {
                while (!frameQueue.isEmpty() && !inputQueue.isEmpty()) {
                    Frame f = frameQueue.poll();
                    int idx = inputQueue.poll();
                    try {
                        ByteBuffer ib = decoder.getInputBuffer(idx);
                        ib.clear();
                        ib.put(f.data);
                        decoder.queueInputBuffer(idx, 0, f.data.length, f.pts, f.flags);
                    } catch (Exception e) {
                        Log.e(TAG, "queue frame failed", e);
                        frameQueue.addFirst(f);
                        inputQueue.addFirst(idx);
                        break;
                    }
                }
            }
        }
    }

    private void releaseDecoder() {
        configured = false;
        synchronized (assembling) {
            lastSeq = -1; // 释放时重置序号，避免重连后首帧被误判为过期
            assembling.clear();
        }
        if (decoder != null) {
            try {
                decoder.stop();
            } catch (Exception ignored) {
            }
            try {
                decoder.release();
            } catch (Exception ignored) {
            }
            decoder = null;
        }
        synchronized (pendingConfigs) {
            pendingConfigs.clear();
        }
        synchronized (frameQueue) {
            frameQueue.clear();
        }
        synchronized (inputQueue) {
            inputQueue.clear();
        }
    }

    public void release() {
        releaseDecoder();
        surface = null;
    }

    public int getVideoWidth() {
        return videoWidth;
    }

    public int getVideoHeight() {
        return videoHeight;
    }

    public void setOnResolutionUpdateListener(OnResolutionUpdateListener l) {
        this.resolutionListener = l;
    }

    /** 分辨率变化监听：当被控端上报新的视频宽高时回调，供 UI 调整渲染视图比例。 */
    public interface OnResolutionUpdateListener {
        void onResolutionUpdated(int width, int height);
    }
}
