package com.anyun.h264;

import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import com.anyun.libusbcamera.UsbCamera;
import com.anyun.libusbcamera.WatermarkParam;
import com.anyun.h264.model.WatermarkInfo;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.concurrent.atomic.AtomicBoolean;

import timber.log.Timber;

/**
 * H264视频编码器
 * 使用UsbCamera获取视频数据，进行H264编码，并通过UDP按JT/T 1076-2016协议上传
 *
 * 使用示例：
 * <pre>
 * // 创建编码器
 * H264Encoder encoder = new H264Encoder();
 *
 * // 设置编码参数
 * encoder.setEncoderParams(640, 480, 25, 2000000);
 *
 * // 设置输出文件（可选，用于保存H264编码数据，可用VLC播放验证）
 * // 使用应用外部存储目录（推荐）
 * File outputFile = new File(context.getExternalFilesDir(null), "test.h264");
 * encoder.setOutputFile(outputFile.getAbsolutePath());
 * encoder.setEnableFileOutput(true); // 启用文件输出，设置为false则不写入文件
 *
 * // 设置UDP服务器地址（可选）
 * encoder.setServerAddress("192.168.1.100", 8888);
 * encoder.setProtocolParams("123456789012", (byte)1);
 * encoder.setEnableNetworkTransmission(true); // 启用TCP/UDP网络传输，false表示禁用
 *
 * // 初始化并启动
 * int[] cameraIdRange = {0, 0};
 * int[] resolution = {640, 480};
 * if (encoder.initialize(cameraIdRange, "camera", resolution, false)) {
 *     encoder.start();
 * }
 *
 * // 停止编码
 * encoder.stop();
 * </pre>
 *
 * 生成的.h264文件可以用VLC播放器直接播放验证。
 */
public class H264Encoder {
    private static final String TAG = "H264Encoder";

    private UsbCamera usbCamera;
    private MediaCodec encoder;
    private Thread encodeThread;
    private AtomicBoolean isRunning = new AtomicBoolean(false);

    // 编码参数
    private int width = 640;
    private int height = 480;
    private int frameRate = 25;
    private int bitrate = 2000000; // 2Mbps
    private int iFrameInterval = 1; // I帧间隔（秒）

    // JT/T 1076-2016 协议工具类
    private JT1076ProtocolHelper protocolHelper;
    private long lastIFrameTime = 0; // 上一个I帧时间
    private long lastFrameTime = 0; // 上一帧时间

    // 文件输出
    private FileOutputStream fileOutputStream;
    private String outputFilePath;
    private boolean enableFileOutput = false; // 是否启用文件输出
    private boolean spsPpsWritten = false; // 标记SPS/PPS是否已写入

    // 网络传输控制
    private boolean enableNetworkTransmission = true; // 是否启用TCP/UDP网络传输

    // 编码回调
    public interface OnFrameEncodedCallback {
        void onFrameEncoded(byte[] data, boolean isKeyFrame);
    }
    private OnFrameEncodedCallback callback;

    public H264Encoder() {
        this.usbCamera = new UsbCamera();
        this.protocolHelper = new JT1076ProtocolHelper();
        protocolHelper.setProtocolType(JT1076ProtocolHelper.PROTOCOL_TYPE_TCP);//设置为tcp传输
    }

    /**
     * 设置编码参数
     */
    public void setEncoderParams(int width, int height, int frameRate, int bitrate) {
        this.width = width;
        this.height = height;
        this.frameRate = frameRate;
        this.bitrate = bitrate;
    }

    /**
     * 设置UDP服务器地址
     */
    public void setServerAddress(String ip, int port) {
        protocolHelper.setServerAddress(ip, port);
    }

    /**
     * 设置SIM卡号和逻辑通道号
     */
    public void setProtocolParams(String simCardNumber, byte logicalChannelNumber) {
        protocolHelper.setProtocolParams(simCardNumber, logicalChannelNumber);
    }

    /**
     * 设置编码回调
     */
    public void setOnFrameEncodedCallback(OnFrameEncodedCallback callback) {
        this.callback = callback;
    }

    /**
     * 设置输出文件路径（用于保存H264编码数据）
     * @param filePath 文件路径，例如："/sdcard/test.h264" 或使用Context.getExternalFilesDir()
     */
    public void setOutputFile(String filePath) {
        this.outputFilePath = filePath;
    }

    /**
     * 设置是否启用文件输出
     * @param enable true表示启用文件输出，false表示禁用
     */
    public void setEnableFileOutput(boolean enable) {
        this.enableFileOutput = enable;
    }

    /**
     * 设置是否启用TCP/UDP网络传输
     * @param enable true表示启用网络传输，false表示禁用
     */
    public void setEnableNetworkTransmission(boolean enable) {
        this.enableNetworkTransmission = enable;
        Timber.d("Network transmission " + (enable ? "enabled" : "disabled"));
    }

    /**
     * 设置水印信息
     * @param watermarkInfo 水印信息对象
     */
    public void setWatermarkInfo(WatermarkInfo watermarkInfo) {
        if (watermarkInfo == null) {
            Timber.w("WatermarkInfo is null, disabling watermark");
            usbCamera.enableWatermark(false, null);
            return;
        }

        try {
            // 构建水印文本列表（分行显示，每行一个信息项）
            ArrayList<WatermarkParam> watermarkParams = new ArrayList<>();
            
            // 从左上角开始，每行间隔25像素
            int yOffset = 30;
            int xOffset = 10;
            
            // 车牌号
            if (watermarkInfo.getPlateNumber() != null && !watermarkInfo.getPlateNumber().isEmpty()) {
                watermarkParams.add(new WatermarkParam(xOffset, yOffset, 
                    "车牌:" + watermarkInfo.getPlateNumber()));
                yOffset += 25;
            }
            
            // 学员姓名
            if (watermarkInfo.getStudent() != null && !watermarkInfo.getStudent().isEmpty()) {
                watermarkParams.add(new WatermarkParam(xOffset, yOffset, 
                    "学员:" + watermarkInfo.getStudent()));
                yOffset += 25;
            }
            
            // 教练姓名
            if (watermarkInfo.getCoach() != null && !watermarkInfo.getCoach().isEmpty()) {
                watermarkParams.add(new WatermarkParam(xOffset, yOffset, 
                    "教练:" + watermarkInfo.getCoach()));
                yOffset += 25;
            }
            
            // 位置信息（纬度,经度）
            if (watermarkInfo.getLongitude() != null && watermarkInfo.getLatitude() != null) {
                watermarkParams.add(new WatermarkParam(xOffset, yOffset, 
                    String.format("位置:%.6f,%.6f", watermarkInfo.getLatitude(), watermarkInfo.getLongitude())));
                yOffset += 25;
            }
            
            // 驾校名称
            if (watermarkInfo.getDrivingSchool() != null && !watermarkInfo.getDrivingSchool().isEmpty()) {
                watermarkParams.add(new WatermarkParam(xOffset, yOffset, 
                    "驾校:" + watermarkInfo.getDrivingSchool()));
                yOffset += 25;
            }
            
            // 车速
            if (watermarkInfo.getSpeed() != null) {
                watermarkParams.add(new WatermarkParam(xOffset, yOffset, 
                    String.format("车速:%.1fkm/h", watermarkInfo.getSpeed())));
            }
            
            if (!watermarkParams.isEmpty()) {
                // 启用水印，使用默认字体路径（如果系统有字体文件）
                // 颜色：0-REVERSE（反色），1-BLACK，2-WHITE，3-RED，4-GREEN，5-BLUE
                // 字体大小、倍数可以根据需要调整
                String fontPath = "/system/fonts/DroidSans.ttf"; // 默认字体路径，如果不存在可以传null
                usbCamera.enableWatermark(true, fontPath);
                
                // 设置水印：颜色（2=白色），字体大小（24），倍数（1），文本列表
                usbCamera.setWatermark(2, 24, 1, watermarkParams);
                
                Timber.d("Watermark set successfully: %s", watermarkInfo);
            } else {
                Timber.w("No watermark text to display, disabling watermark");
                usbCamera.enableWatermark(false, null);
            }
        } catch (Exception e) {
            Timber.e(e, "Failed to set watermark");
            usbCamera.enableWatermark(false, null);
        }
    }

    /**
     * 初始化摄像头和编码器
     */
    public boolean initialize(int[] cameraIdRange, String cameraName, int[] resolution, boolean ayCamera) {
        try {
            // 1. setenv
            usbCamera.setenv();

            // 2. prepareCamera (最多尝试3次：初始1次 + 重试2次)
            int[] actualResolution = new int[2];
            actualResolution[0] = 640;
            actualResolution[1] = 480;
            System.arraycopy(resolution, 0, actualResolution, 0, 2);

            int result = -1;
            int maxRetries = 3; // 总共尝试3次
            for (int attempt = 0; attempt < maxRetries; attempt++) {
                result = usbCamera.prepareCamera(cameraIdRange, cameraName, actualResolution, ayCamera);
                if (result == 0) {
                    // 成功，跳出循环
                    if (attempt > 0) {
                        Timber.d("prepareCamera succeeded on attempt " + (attempt + 1));
                    }
                    break;
                } else {
                    // 失败，记录日志
                    Timber.w( "prepareCamera failed on attempt " + (attempt + 1) + ": " + result);
                    if (attempt < maxRetries - 1) {
                        Timber.d( "Retrying prepareCamera...");
                    }
                }
            }

            if (result != 0) {
                Timber.e("prepareCamera failed after " + maxRetries + " attempts: " + result);
                return false;
            }

            // 更新实际分辨率
            width = actualResolution[0];
            height = actualResolution[1];
            Timber.d("Camera initialized with resolution: " + width + "x" + height);

            // 3. 初始化H264编码器
            initEncoder();

            // 4. 初始化Socket（UDP或TCP，根据协议类型自动选择）
            // 只有在启用网络传输时才初始化Socket
            if (enableNetworkTransmission) {
                if (!protocolHelper.initializeSocket()) {
                    return false;
                }
            } else {
                Timber.d("Network transmission disabled, skipping socket initialization");
            }

            // 5. 初始化文件输出（仅创建文件，SPS/PPS在第一次输出时写入）
            if (enableFileOutput && outputFilePath != null && !outputFilePath.isEmpty()) {
                if (!initFileOutput()) {
                    Timber.w("File output initialization failed, continuing without file output");
                }
            }

            return true;
        } catch (Exception e) {
            Timber.e(e,"Initialize failed");
            return false;
        }
    }

    /**
     * 初始化H264编码器
     */
    private void initEncoder() throws IOException {
        MediaFormat format = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC, width, height);
        format.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Flexible);
        format.setInteger(MediaFormat.KEY_BIT_RATE, bitrate);
        format.setInteger(MediaFormat.KEY_FRAME_RATE, frameRate);
        format.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, iFrameInterval);

        encoder = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_AVC);
        encoder.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
        encoder.start();

        Timber.d( "H264 encoder initialized");
    }

    /**
     * 初始化文件输出
     * @return 是否成功
     */
    private boolean initFileOutput() {
        try {
            File file = new File(outputFilePath);
            File parentDir = file.getParentFile();
            if (parentDir != null && !parentDir.exists()) {
                boolean created = parentDir.mkdirs();
                if (!created && !parentDir.exists()) {
                    Timber.e("Failed to create parent directory: " + parentDir.getAbsolutePath());
                    return false;
                }
            }

            fileOutputStream = new FileOutputStream(file);
            spsPpsWritten = false;

            Timber.d("File output initialized: " + outputFilePath);
            return true;
        } catch (Exception e) {
            Timber.e(e,"Initialize file output failed");
            if (fileOutputStream != null) {
                try {
                    fileOutputStream.close();
                } catch (IOException ie) {
                    Timber.e(ie, "Close file output stream failed");
                }
                fileOutputStream = null;
            }
            return false;
        }
    }

    /**
     * 写入SPS/PPS到文件（从CSD或关键帧数据中提取）
     */
    private void writeSpsPpsToFile() {
        if (!enableFileOutput || fileOutputStream == null || spsPpsWritten) {
            return;
        }

        try {
            // 尝试从编码器输出格式中获取CSD
            MediaFormat format = encoder.getOutputFormat();
            ByteBuffer spsBuffer = format.getByteBuffer("csd-0"); // SPS
            ByteBuffer ppsBuffer = format.getByteBuffer("csd-1"); // PPS

            if (spsBuffer != null && ppsBuffer != null) {
                // CSD格式通常是AVCC格式，需要转换为Annex-B
                byte[] sps = new byte[spsBuffer.remaining()];
                byte[] pps = new byte[ppsBuffer.remaining()];
                spsBuffer.get(sps);
                ppsBuffer.get(pps);

                // 写入SPS和PPS到文件（Annex-B格式）
                byte[] nalStartCode = {0x00, 0x00, 0x00, 0x01};

                // CSD数据通常是AVCC格式：前4字节是大端格式的长度，后面是NAL数据
                // 检查并跳过长度前缀
                int spsOffset = 0;
                int ppsOffset = 0;
                int spsLength = sps.length;
                int ppsLength = pps.length;

                // 检查是否为AVCC格式（前4字节是长度，通常不会是0x00000001）
                if (sps.length > 4 && (sps[0] != 0x00 || sps[1] != 0x00 || sps[2] != 0x00 || sps[3] != 0x01)) {
                    // AVCC格式：前4字节是长度（大端）
                    spsOffset = 4;
                    spsLength = sps.length - 4;
                }
                if (pps.length > 4 && (pps[0] != 0x00 || pps[1] != 0x00 || pps[2] != 0x00 || pps[3] != 0x01)) {
                    // AVCC格式：前4字节是长度（大端）
                    ppsOffset = 4;
                    ppsLength = pps.length - 4;
                }

                // 写入SPS
                fileOutputStream.write(nalStartCode);
                fileOutputStream.write(sps, spsOffset, spsLength);

                // 写入PPS
                fileOutputStream.write(nalStartCode);
                fileOutputStream.write(pps, ppsOffset, ppsLength);
                fileOutputStream.flush();

                spsPpsWritten = true;
                Timber.d("SPS/PPS written to file, SPS size: " + spsLength + ", PPS size: " + ppsLength);
            } else {
                Timber.w("SPS/PPS not found in CSD, will extract from first key frame");
            }
        } catch (Exception e) {
            Timber.e(e,"Write SPS/PPS to file error");
        }
    }

    /**
     * 开始编码
     */
    public void start() {
        if (isRunning.get()) {
            Timber.w("Encoder is already running");
            return;
        }

        isRunning.set(true);

        // 启动编码线程
        encodeThread = new Thread(new Runnable() {
            @Override
            public void run() {
                encodeLoop();
            }
        });
        encodeThread.start();

        Timber.d("H264 encoder started");
    }

    /**
     * 编码循环
     */
    private void encodeLoop() {
        // YUV420P格式: Y平面 + U平面 + V平面
        // Y平面大小 = width * height
        // U平面大小 = width * height / 4
        // V平面大小 = width * height / 4
        // 总大小 = width * height * 3 / 2
        byte[] yuvBuffer = new byte[width * height * 3 / 2];

        MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();

        while (isRunning.get()) {
            try {
                // processCamera - 读取一帧
                int processResult = usbCamera.processCamera();
                if (processResult != 0) {
                    Timber.w("processCamera returned: " + processResult);
                    Thread.sleep(10);
                    continue;
                }

                // 获取RGBA数据 (type=1 表示推流，输出YUV420P格式)
                usbCamera.rgba(0, yuvBuffer);

                // 将YUV420P数据送入编码器
                long timestamp = System.currentTimeMillis();
                encodeFrame(nv12ToNV21(yuvBuffer,width,height), timestamp, bufferInfo);

            } catch (Exception e) {
                Timber.e(e, "Encode loop error");
                try {
                    Thread.sleep(100);
                } catch (InterruptedException ie) {
                    break;
                }
            }
        }

        Timber.d( "Encode loop exited");
    }

    /**
     * 编码一帧数据
     */
    private void encodeFrame(byte[] yuvData, long timestamp, MediaCodec.BufferInfo bufferInfo) {
        try {
            // 获取输入缓冲区
            int inputBufferIndex = encoder.dequeueInputBuffer(10000);
            if (inputBufferIndex >= 0) {
                ByteBuffer inputBuffer = encoder.getInputBuffer(inputBufferIndex);
                if (inputBuffer != null) {
                    inputBuffer.clear();
                    inputBuffer.put(yuvData);
                    encoder.queueInputBuffer(inputBufferIndex, 0, yuvData.length, timestamp * 1000, 0);
                }
            }

            // 获取输出数据
            int outputBufferIndex = encoder.dequeueOutputBuffer(bufferInfo, 0);
            while (outputBufferIndex >= 0) {
                ByteBuffer outputBuffer = encoder.getOutputBuffer(outputBufferIndex);
                if (outputBuffer != null && bufferInfo.size > 0) {
                    // 检查是否为关键帧
                    boolean isKeyFrame = (bufferInfo.flags & MediaCodec.BUFFER_FLAG_KEY_FRAME) != 0;
                    // 复制编码数据
                    byte[] encodedData = new byte[bufferInfo.size];
                    outputBuffer.position(bufferInfo.offset);
                    outputBuffer.get(encodedData, 0, bufferInfo.size);

                    // 写入文件
                    writeToFile(encodedData, isKeyFrame);

                    // 发送编码数据
                    sendEncodedData(encodedData, timestamp, isKeyFrame);

                    // 回调
                    if (callback != null) {
                        callback.onFrameEncoded(encodedData, isKeyFrame);
                    }
                }

                encoder.releaseOutputBuffer(outputBufferIndex, false);
                outputBufferIndex = encoder.dequeueOutputBuffer(bufferInfo, 0);
            }

        } catch (Exception e) {
            Timber.e(e,"Encode frame error");
        }
    }

    /**
     * 将编码数据写入H264文件
     */
    private void writeToFile(byte[] data, boolean isKeyFrame) {
        if (!enableFileOutput || fileOutputStream == null) {
            return;
        }

        try {
            // 如果是第一个关键帧，确保SPS/PPS已写入
            if (isKeyFrame && !spsPpsWritten) {
                writeSpsPpsToFile();

                // 如果从CSD获取失败，尝试从关键帧数据中提取
                // MediaCodec输出的关键帧通常已经包含SPS/PPS，但为了确保文件完整性，
                // 我们已经从CSD写入，这里直接写入关键帧数据即可
                if (!spsPpsWritten) {
                    Timber.d("SPS/PPS will be included in key frame data");
                }
            }

            // MediaCodec输出的H264数据已经是Annex-B格式（包含0x00000001分隔符）
            // 直接写入文件即可
            fileOutputStream.write(data);
            fileOutputStream.flush();
        } catch (IOException e) {
            Timber.e(e, "Write to file error");
        }
    }

    /**
     * 发送编码后的数据（按JT/T 1076-2016协议打包）
     */
    private void sendEncodedData(byte[] data, long timestamp, boolean isKeyFrame) {
        // 如果未启用网络传输，直接返回
        if (!enableNetworkTransmission) {
            return;
        }

        try {
            // 计算时间间隔
            long currentTime = System.currentTimeMillis();
            long lastIFrameInterval = (lastIFrameTime > 0) ? (currentTime - lastIFrameTime) : 0;
            long lastFrameInterval = (lastFrameTime > 0) ? (currentTime - lastFrameTime) : 0;

            if (isKeyFrame) {
                lastIFrameTime = currentTime;
            }
            lastFrameTime = currentTime;

            // 判断帧类型
            int dataType = isKeyFrame ? JT1076ProtocolHelper.DATA_TYPE_I_FRAME :
                    JT1076ProtocolHelper.DATA_TYPE_P_FRAME;

            // 分包发送（如果数据超过最大包大小）
            int offset = 0;
            int totalPackets = (int) Math.ceil((double) data.length / JT1076ProtocolHelper.MAX_PACKET_SIZE);

            for (int i = 0; i < totalPackets; i++) {
                int packetDataSize = Math.min(JT1076ProtocolHelper.MAX_PACKET_SIZE, data.length - offset);
                byte[] packetData = Arrays.copyOfRange(data, offset, offset + packetDataSize);

                // 确定分包标记
                int packetMark;
                if (totalPackets == 1) {
                    packetMark = JT1076ProtocolHelper.PACKET_MARK_ATOMIC;
                } else if (i == 0) {
                    packetMark = JT1076ProtocolHelper.PACKET_MARK_FIRST;
                } else if (i == totalPackets - 1) {
                    packetMark = JT1076ProtocolHelper.PACKET_MARK_LAST;
                } else {
                    packetMark = JT1076ProtocolHelper.PACKET_MARK_MIDDLE;
                }

                // 创建RTP包
                byte[] rtpPacket = protocolHelper.createVideoRtpPacket(
                        packetData, timestamp, dataType, packetMark,
                        lastIFrameInterval, lastFrameInterval);

                // 发送RTP包（UDP或TCP，根据协议类型自动选择）
                protocolHelper.sendPacket(rtpPacket);

                offset += packetDataSize;
            }

        } catch (Exception e) {
            Timber.e(e,"Send encoded data error");
        }
    }

    /**
     * 停止编码
     */
    public void stop() {
        if (!isRunning.get()) {
            return;
        }

        isRunning.set(false);

        // 等待编码线程结束
        if (encodeThread != null) {
            try {
                encodeThread.join(2000);
            } catch (InterruptedException e) {
                Timber.e(e, "Wait encode thread error");
            }
        }

        // 停止摄像头
        if (usbCamera != null) {
            usbCamera.stopCamera();
        }

        // 释放编码器
        if (encoder != null) {
            try {
                encoder.stop();
                encoder.release();
                encoder = null;
            } catch (Exception e) {
                Timber.e(e, "Release encoder error");
            }
        }

        // 关闭Socket（UDP或TCP，根据协议类型自动选择）
        // 只有在启用网络传输时才需要关闭Socket
        if (enableNetworkTransmission && protocolHelper != null) {
            protocolHelper.closeSocket();
        }

        // 关闭文件输出
        closeFileOutput();

        Timber.d("H264 encoder stopped");
    }

    /**
     * 关闭文件输出
     */
    private void closeFileOutput() {
        if (fileOutputStream != null) {
            try {
                fileOutputStream.flush();
                fileOutputStream.close();
                Timber.d("File output closed: " + outputFilePath);
            } catch (IOException e) {
                Timber.e(e, "Close file output error");
            } finally {
                fileOutputStream = null;
                spsPpsWritten = false;
            }
        }
    }

    /**
     * 释放资源
     */
    public void release() {
        stop();
    }

    byte[] ret = null;
    // YYYYYYYY UVUV(nv12)--> YYYYYYYY VUVU(nv21)
    private byte[] nv12ToNV21(byte[] nv12, int width, int height) {
//        Log.i(TAG,"nv12toNv21:"+width+"height:"+height);
        if (ret == null){
            ret = new byte[width * height * 3 /2];
        }
        int framesize = width * height;
        int i = 0, j = 0;
        // 拷贝Y分量
        System.arraycopy(nv12, 0,ret , 0, framesize);
        // 拷贝UV分量
        for (j = framesize; j < nv12.length; j += 2) {
            ret[j] = nv12[j+1];
            ret[j+1] = nv12[j];
        }
        return ret;
    }
}