package cn.iocoder.byzs.module.ai.framework.ai.core.model.speech;

import com.alibaba.nls.client.AccessToken;
import com.alibaba.nls.client.protocol.NlsClient;
import com.alibaba.nls.client.protocol.InputFormatEnum;
import com.alibaba.nls.client.protocol.SampleRateEnum;
import com.alibaba.nls.client.protocol.asr.SpeechTranscriber;
import com.alibaba.nls.client.protocol.asr.SpeechTranscriberListener;
import com.alibaba.nls.client.protocol.asr.SpeechTranscriberResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.sound.sampled.AudioFormat;
import javax.sound.sampled.AudioSystem;
import javax.sound.sampled.DataLine;
import javax.sound.sampled.TargetDataLine;

public class RealTimeSpeechTranscriberDemo {
    private static final Logger logger = LoggerFactory.getLogger(RealTimeSpeechTranscriberDemo.class);
    private String appKey;
    private NlsClient client;

    public RealTimeSpeechTranscriberDemo(String appKey, String id, String secret, String url) {
        this.appKey = appKey;
        //创建NlsClient实例，应用全局创建一个即可，默认服务地址为阿里云线上服务地址。
        //获取token，实际使用时注意在accessToken.getExpireTime()过期前再次获取。
        AccessToken accessToken = new AccessToken(id, secret);
        try {
            accessToken.apply();
            System.out.println("get token: " + ", expire time: " + accessToken.getExpireTime());
            if(url.isEmpty()) {
                client = new NlsClient(accessToken.getToken());
            }else {
                client = new NlsClient(url, accessToken.getToken());
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private SpeechTranscriberListener getTranscriberListener() {
        return new SpeechTranscriberListener() {
            //识别出中间结果。仅当setEnableIntermediateResult为true时，才会返回该消息。
            @Override
            public void onTranscriptionResultChange(SpeechTranscriberResponse response) {
                System.out.println("task_id: " + response.getTaskId() +
                    ", name: " + response.getName() +
                    //状态码"20000000"表示正常识别。
                    ", status: " + response.getStatus() +
                    //句子编号，从1开始递增。
                    ", index: " + response.getTransSentenceIndex() +
                    //当前的识别结果。
                    ", result: " + response.getTransSentenceText() +
                    //当前已处理的音频时长，单位为毫秒。
                    ", time: " + response.getTransSentenceTime());
            }

            @Override
            public void onTranscriberStart(SpeechTranscriberResponse response) {
                //task_id是调用方和服务端通信的唯一标识，遇到问题时，需要提供此task_id。
                System.out.println("task_id: " + response.getTaskId() + ", name: " + response.getName() + ", status: " + response.getStatus());
            }

            @Override
            public void onSentenceBegin(SpeechTranscriberResponse response) {
                System.out.println("task_id: " + response.getTaskId() + ", name: " + response.getName() + ", status: " + response.getStatus());

            }

            //识别出一句话。服务端会智能断句，当识别到一句话结束时会返回此消息。
            @Override
            public void onSentenceEnd(SpeechTranscriberResponse response) {
                System.out.println("task_id: " + response.getTaskId() +
                    ", name: " + response.getName() +
                    //状态码"20000000"表示正常识别。
                    ", status: " + response.getStatus() +
                    //句子编号，从1开始递增。
                    ", index: " + response.getTransSentenceIndex() +
                    //当前的识别结果。
                    ", result: " + response.getTransSentenceText() +
                    //置信度
                    ", confidence: " + response.getConfidence() +
                    //开始时间
                    ", begin_time: " + response.getSentenceBeginTime() +
                    //当前已处理的音频时长，单位为毫秒。
                    ", time: " + response.getTransSentenceTime());
            }

            //识别完毕
            @Override
            public void onTranscriptionComplete(SpeechTranscriberResponse response) {
                System.out.println("task_id: " + response.getTaskId() + ", name: " + response.getName() + ", status: " + response.getStatus());
            }

            @Override
            public void onFail(SpeechTranscriberResponse response) {
                //task_id是调用方和服务端通信的唯一标识，遇到问题时，需要提供此task_id。
                System.out.println("task_id: " + response.getTaskId() +  ", status: " + response.getStatus() + ", status_text: " + response.getStatusText());
            }
        };
    }

    public void process() {
        SpeechTranscriber transcriber = null;
        TargetDataLine targetDataLine = null;
        try {
            //创建实例、建立连接。
            transcriber = new SpeechTranscriber(client, getTranscriberListener());
            transcriber.setAppKey(appKey);
            //输入音频编码方式。
            transcriber.setFormat(InputFormatEnum.PCM);
            //输入音频采样率。
            transcriber.setSampleRate(SampleRateEnum.SAMPLE_RATE_16K);
            //是否返回中间识别结果。
            transcriber.setEnableIntermediateResult(false);
            //是否生成并返回标点符号。
            transcriber.setEnablePunctuation(true);
            //是否将返回结果规整化，比如将一百返回为100。
            transcriber.setEnableITN(false);

            //此方法将以上参数设置序列化为JSON发送给服务端，并等待服务端确认。
            transcriber.start();

            //配置音频采集格式（16kHz，16位，单声道，小端序）
            AudioFormat format = new AudioFormat(16000, 16, 1, true, false);
            DataLine.Info info = new DataLine.Info(TargetDataLine.class, format);
            
            //打开麦克风输入流
            targetDataLine = (TargetDataLine) AudioSystem.getLine(info);
            targetDataLine.open(format);
            targetDataLine.start();

            //缓冲区大小（3200字节 = 100ms音频）
            byte[] buffer = new byte[3200];
            int len;

            System.out.println("开始实时采集音频...");
            System.out.println("按Ctrl+C停止程序");
            
            //实时读取音频数据并发送
            while (true) {
                len = targetDataLine.read(buffer, 0, buffer.length);
                if (len > 0) {
                    logger.info("send data pack length: " + len);
                    transcriber.send(buffer, len);
                }
            }
        } catch (Exception e) {
            System.err.println(e.getMessage());
        } finally {
            //停止音频采集
            if (targetDataLine != null) {
                targetDataLine.stop();
                targetDataLine.close();
            }
            //通知服务端语音数据发送完毕，等待服务端处理完成。
            if (transcriber != null) {
                try {
                    long now = System.currentTimeMillis();
                    logger.info("ASR wait for complete");
                    transcriber.stop();
                    logger.info("ASR latency : " + (System.currentTimeMillis() - now) + " ms");
                } catch (Exception e) {
                    e.printStackTrace();
                } finally {
                    transcriber.close();
                }
            }
        }
    }

    public void shutdown() {
        client.shutdown();
    }

    public static void main(String[] args) throws Exception {
        String appKey = "4SUOF4LfaU7FekyW";
        String id = "LTAI5tQhMPLXtSgXiPiWbw6D";
        String secret = "HCXpFYjl4swk0qwfIKa9s2bXx0AWcG";
        String url = System.getenv().getOrDefault("NLS_GATEWAY_URL", "wss://nls-gateway-cn-shanghai.aliyuncs.com/ws/v1");
        
        System.out.println("实时语音识别Demo启动中...");
        System.out.println("请确保已正确配置AccessKey凭证");
        
        RealTimeSpeechTranscriberDemo demo = new RealTimeSpeechTranscriberDemo(appKey, id, secret, url);
        demo.process();
        demo.shutdown();
    }
}