package com.anyun.h264;

import android.media.AudioFormat;
import android.media.AudioRecord;
import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import android.util.Log;

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

/**
 * AAC音频编码器
 * 采集音频数据，进行AAC编码，并通过UDP按JT/T 1076-2016协议上传
 */
public class AACEncoder {
    private static final String TAG = "AACEncoder";
    
    private MediaCodec encoder;
    private AudioRecord audioRecord;
    private Thread encodeThread;
    private AtomicBoolean isRunning = new AtomicBoolean(false);
    
    // 音频参数
    private int sampleRate = 16000; // 采样率
    private int channelCount = 1; // 声道数（1=单声道，2=立体声）
    private int bitrate = 64000; // 比特率
    private int audioFormat = AudioFormat.ENCODING_PCM_16BIT;
    
    // JT/T 1076-2016 协议工具类
    private JT1076ProtocolHelper protocolHelper;
    private long lastFrameTime = 0; // 上一帧时间
    
    // 编码回调
    public interface OnFrameEncodedCallback {
        void onFrameEncoded(byte[] data);
    }
    private OnFrameEncodedCallback callback;
    
    public AACEncoder() {
        this.protocolHelper = new JT1076ProtocolHelper();
    }
    
    /**
     * 设置音频参数
     */
    public void setAudioParams(int sampleRate, int channelCount, int bitrate) {
        this.sampleRate = sampleRate;
        this.channelCount = channelCount;
        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;
    }
    
    /**
     * 初始化音频录制和编码器
     */
    public boolean initialize() {
        try {
            // 1. 初始化AudioRecord
            int channelConfig = channelCount == 1 ? 
                AudioFormat.CHANNEL_IN_MONO : AudioFormat.CHANNEL_IN_STEREO;
            int bufferSize = AudioRecord.getMinBufferSize(sampleRate, channelConfig, audioFormat);
            
            if (bufferSize == AudioRecord.ERROR_BAD_VALUE || bufferSize == AudioRecord.ERROR) {
                Log.e(TAG, "Invalid audio parameters");
                return false;
            }
            
            // 使用更大的缓冲区以避免欠载
            bufferSize *= 4;
            
            audioRecord = new AudioRecord(
                android.media.MediaRecorder.AudioSource.MIC,
                sampleRate,
                channelConfig,
                audioFormat,
                bufferSize
            );
            
            if (audioRecord.getState() != AudioRecord.STATE_INITIALIZED) {
                Log.e(TAG, "AudioRecord initialization failed");
                return false;
            }
            
            Log.d(TAG, "AudioRecord initialized: sampleRate=" + sampleRate + 
                  ", channels=" + channelCount + ", bufferSize=" + bufferSize);
            
            // 2. 初始化AAC编码器
            initEncoder();
            
            // 3. 初始化Socket（UDP或TCP，根据协议类型自动选择）
            if (!protocolHelper.initializeSocket()) {
                return false;
            }
            
            return true;
        } catch (Exception e) {
            Log.e(TAG, "Initialize failed", e);
            return false;
        }
    }
    
    /**
     * 初始化AAC编码器
     */
    private void initEncoder() throws IOException {
        MediaFormat format = MediaFormat.createAudioFormat(MediaFormat.MIMETYPE_AUDIO_AAC, 
                                                          sampleRate, channelCount);
        format.setInteger(MediaFormat.KEY_AAC_PROFILE, MediaCodecInfo.CodecProfileLevel.AACObjectLC);
        format.setInteger(MediaFormat.KEY_BIT_RATE, bitrate);
        format.setInteger(MediaFormat.KEY_MAX_INPUT_SIZE, 4096);
        
        encoder = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_AUDIO_AAC);
        encoder.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
        encoder.start();
        
        Log.d(TAG, "AAC encoder initialized");
    }
    
    /**
     * 开始编码
     */
    public void start() {
        if (isRunning.get()) {
            Log.w(TAG, "Encoder is already running");
            return;
        }
        
        if (audioRecord == null || encoder == null) {
            Log.e(TAG, "Encoder not initialized");
            return;
        }
        
        isRunning.set(true);
        
        // 开始录音
        audioRecord.startRecording();
        
        // 启动编码线程
        encodeThread = new Thread(new Runnable() {
            @Override
            public void run() {
                encodeLoop();
            }
        });
        encodeThread.start();
        
        Log.d(TAG, "AAC encoder started");
    }
    
    /**
     * 编码循环
     */
    private void encodeLoop() {
        // 读取缓冲区大小（1024个采样点）
        int inputBufferSize = sampleRate * channelCount * 2 / 25; // 40ms的音频数据
        byte[] pcmBuffer = new byte[inputBufferSize];
        
        MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
        
        while (isRunning.get()) {
            try {
                // 从AudioRecord读取PCM数据
                int bytesRead = audioRecord.read(pcmBuffer, 0, pcmBuffer.length);
                
                if (bytesRead < 0) {
                    Log.w(TAG, "AudioRecord read error: " + bytesRead);
                    Thread.sleep(10);
                    continue;
                }
                
                // 将PCM数据送入编码器
                long timestamp = System.currentTimeMillis();
                encodeFrame(pcmBuffer, bytesRead, timestamp, bufferInfo);
                
            } catch (Exception e) {
                Log.e(TAG, "Encode loop error", e);
                try {
                    Thread.sleep(10);
                } catch (InterruptedException ie) {
                    break;
                }
            }
        }
        
        Log.d(TAG, "Encode loop exited");
    }
    
    /**
     * 编码一帧数据
     */
    private void encodeFrame(byte[] pcmData, int dataSize, 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(pcmData, 0, dataSize);
                    encoder.queueInputBuffer(inputBufferIndex, 0, dataSize, timestamp * 1000, 0);
                }
            }
            
            // 获取输出数据
            int outputBufferIndex = encoder.dequeueOutputBuffer(bufferInfo, 0);
            while (outputBufferIndex >= 0) {
                ByteBuffer outputBuffer = encoder.getOutputBuffer(outputBufferIndex);
                if (outputBuffer != null && bufferInfo.size > 0) {
                    // 复制编码数据
                    byte[] encodedData = new byte[bufferInfo.size];
                    outputBuffer.position(bufferInfo.offset);
                    outputBuffer.get(encodedData, 0, bufferInfo.size);
                    
                    // 发送编码数据
                    sendEncodedData(encodedData, timestamp);
                    
                    // 回调
                    if (callback != null) {
                        callback.onFrameEncoded(encodedData);
                    }
                }
                
                encoder.releaseOutputBuffer(outputBufferIndex, false);
                outputBufferIndex = encoder.dequeueOutputBuffer(bufferInfo, 0);
            }
            
        } catch (Exception e) {
            Log.e(TAG, "Encode frame error", e);
        }
    }
    
    /**
     * 发送编码后的数据（按JT/T 1076-2016协议打包）
     */
    private void sendEncodedData(byte[] data, long timestamp) {
        try {
            // 计算时间间隔
            long currentTime = System.currentTimeMillis();
            long lastFrameInterval = (lastFrameTime > 0) ? (currentTime - lastFrameTime) : 0;
            lastFrameTime = currentTime;
            
            // 分包发送（如果数据超过最大包大小）
            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包（音频不需要Last I Frame Interval和Last Frame Interval字段）
                byte[] rtpPacket = protocolHelper.createAudioRtpPacket(
                    packetData, timestamp, JT1076ProtocolHelper.DATA_TYPE_AUDIO, packetMark);
                
                // 发送RTP包（UDP或TCP，根据协议类型自动选择）
                protocolHelper.sendPacket(rtpPacket);
                
                offset += packetDataSize;
            }
            
        } catch (Exception e) {
            Log.e(TAG, "Send encoded data error", e);
        }
    }
    
    /**
     * 停止编码
     */
    public void stop() {
        if (!isRunning.get()) {
            return;
        }
        
        isRunning.set(false);
        
        // 停止录音
        if (audioRecord != null) {
            try {
                audioRecord.stop();
            } catch (Exception e) {
                Log.e(TAG, "Stop AudioRecord error", e);
            }
        }
        
        // 等待编码线程结束
        if (encodeThread != null) {
            try {
                encodeThread.join(2000);
            } catch (InterruptedException e) {
                Log.e(TAG, "Wait encode thread error", e);
            }
        }
        
        // 释放编码器
        if (encoder != null) {
            try {
                encoder.stop();
                encoder.release();
                encoder = null;
            } catch (Exception e) {
                Log.e(TAG, "Release encoder error", e);
            }
        }
        
        // 释放AudioRecord
        if (audioRecord != null) {
            audioRecord.release();
            audioRecord = null;
        }
        
        // 关闭Socket（UDP或TCP，根据协议类型自动选择）
        if (protocolHelper != null) {
            protocolHelper.closeSocket();
        }
        
        Log.d(TAG, "AAC encoder stopped");
    }
    
    /**
     * 释放资源
     */
    public void release() {
        stop();
    }
}