package com.rj.speech;

import com.alibaba.dashscope.audio.asr.transcription.*;
import com.alibaba.dashscope.common.TaskStatus;
import com.google.gson.*;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.net.HttpURLConnection;
import java.net.URL;
import java.util.*;

/**
 * 多人对话语音识别测试 - 模型优化版
 * 实现识别音频中的多人会话，并在文字前加上说话人标记
 * Author: Li Zhonghua  wx: spllzh  email(qq): 28668817@qq.com
 * Date: 2025/8/22
 */
public class MultiSpeakerTranscriptionTest {
    
    public static void main(String[] args) {
        String apiKey = System.getenv("DASHSCOPE_API_KEY");
        if (apiKey == null || apiKey.trim().isEmpty()) {
            System.err.println("错误: 请设置环境变量 DASHSCOPE_API_KEY");
            System.exit(1);
        }

        System.out.println("开始测试不同模型的多人对话识别效果...");
        
        // 测试不同模型
        try {
            // 测试推荐的模型
//            testModel(apiKey, "paraformer-v2", 2);
            // 增加fun-asr模型测试
//            testModel(apiKey, "fun-asr", 2);
            // 增加qwen-audio-asr模型测试
            testModel(apiKey, "qwen-audio-asr", 2);
        } catch (Exception e) {
            System.err.println("测试过程中发生错误: " + e.getMessage());
            e.printStackTrace();
        }
    }
    
    /**
     * 测试指定模型的识别效果
     */
    private static void testModel(String apiKey, String model, int speakerCount) {
        System.out.println("\n==========================================");
        System.out.println("测试模型: " + model);
        System.out.println("==========================================");
        
        try {
            TranscriptionParam param = null;
            
            // 根据模型类型创建不同的参数
            if (model.equals("paraformer-v2")) {
                param = TranscriptionParam.builder()
                    .apiKey(apiKey)
                    .model(model)
                    .fileUrls(Arrays.asList(
                            "https://api.huayang-star.com/api/audio/1957424645898579969_a2db5d76-b52c-4289-9c7e-ed55d7f15e42.mp3"
                    ))
                    .parameter("language_hints", new String[]{"zh", "en"})
                    .parameter("enable_speaker_diarization", true)
                    .parameter("speaker_count", speakerCount)
                    .parameter("min_speaker_duration", 0.3)
                    .parameter("noise_level", "low")
                    .parameter("domain", "general")
                    .parameter("speaker_diarization_confidence_threshold", 0.1)
                    .parameter("max_speaker_count", 3)
                    .parameter("enable_word_time_offsets", true)
                    .build();
            } else if (model.equals("fun-asr")) {
                // 为fun-asr模型配置参数
                param = TranscriptionParam.builder()
                    .apiKey(apiKey)
                    .model(model)
                    .fileUrls(Arrays.asList(
                            "https://api.huayang-star.com/api/audio/1957424645898579969_a2db5d76-b52c-4289-9c7e-ed55d7f15e42.mp3"
                    ))
                    .parameter("enable_speaker_diarization", true)
                    .parameter("speaker_count", speakerCount)
                    .parameter("min_speaker_duration", 0.2)
                    .parameter("speaker_diarization_confidence_threshold", 0.05)
                    .parameter("max_speaker_count", 3)
                    .build();
            } else if (model.equals("qwen-audio-asr")) {
                // 为qwen-audio-asr模型配置参数
                param = TranscriptionParam.builder()
                    .apiKey(apiKey)
                    .model(model)
                    .fileUrls(Arrays.asList(
                            "https://api.huayang-star.com/api/audio/1957424645898579969_a2db5d76-b52c-4289-9c7e-ed55d7f15e42.mp3"
                    ))
                    .parameter("enable_speaker_diarization", true)
                    .parameter("speaker_count", speakerCount)
                    .parameter("min_speaker_duration", 0.25)
                    .parameter("speaker_diarization_confidence_threshold", 0.1)
                    .parameter("max_speaker_count", 3)
                    .build();
            } else {
                // 其他模型的基础配置
                param = TranscriptionParam.builder()
                    .apiKey(apiKey)
                    .model(model)
                    .fileUrls(Arrays.asList(
                            "https://api.huayang-star.com/api/audio/1957424645898579969_a2db5d76-b52c-4289-9c7e-ed55d7f15e42.mp3"
                    ))
                    .build();
            }
            
            Transcription transcription = new Transcription();
            
            System.out.println("提交识别请求，使用模型: " + model);
            TranscriptionResult result = transcription.asyncCall(param);
            System.out.println("RequestId: " + result.getRequestId());
            System.out.println("TaskId: " + result.getTaskId());
            
            // 轮询任务状态
            System.out.println("轮询任务状态...");
            int maxRetries = 60;
            int retryCount = 0;
            
            while (retryCount < maxRetries) {
                result = transcription.fetch(TranscriptionQueryParam.FromTranscriptionParam(param, result.getTaskId()));
                System.out.println("任务状态: " + result.getTaskStatus() + " (重试次数: " + (retryCount + 1) + ")");
                
                if (result.getTaskStatus() == TaskStatus.SUCCEEDED) {
                    System.out.println(model + " 识别任务成功完成!");
                    System.out.println("------------------------------------------");
                    
                    // 处理识别结果
                    handleTranscriptionResults(result, model);
                    break;
                } else if (result.getTaskStatus() == TaskStatus.FAILED) {
                    System.err.println(model + " 任务失败");
                    System.err.println("失败原因: " + result.getOutput());
                    break;
                }
                
                Thread.sleep(5000);
                retryCount++;
            }
            
            if (retryCount >= maxRetries) {
                System.err.println(model + " 任务超时");
            }
            
        } catch (Exception e) {
            System.err.println(model + " 发生错误: " + e.getMessage());
            e.printStackTrace();
        }
    }
    
    /**
     * 处理识别结果
     * @param result 识别结果
     * @param model 使用的模型
     */
    private static void handleTranscriptionResults(TranscriptionResult result, String model) {
        try {
            // 获取任务结果列表
            List<TranscriptionTaskResult> taskResultList = result.getResults();
            if (taskResultList != null && !taskResultList.isEmpty()) {
                System.out.println("找到 " + taskResultList.size() + " 个识别结果文件");
                
                for (int i = 0; i < taskResultList.size(); i++) {
                    TranscriptionTaskResult taskResult = taskResultList.get(i);
                    String transcriptionUrl = taskResult.getTranscriptionUrl();
                    
                    System.out.println("\n=== 处理第 " + (i + 1) + " 个识别结果 ===");
                    System.out.println("模型: " + model);
                    System.out.println("下载URL: " + transcriptionUrl);
                    
                    // 下载识别结果
                    String jsonContent = downloadTranscriptionResult(transcriptionUrl);
                    if (jsonContent != null) {
                        // 解析并输出识别内容
                        parseAndOutputTranscription(jsonContent, model);
                    }
                }
            } else {
                // 尝试直接使用result.getOutput()
                System.out.println("\n尝试直接解析输出结果...");
                Object output = result.getOutput();
                if (output != null) {
                    System.out.println("输出结果类型: " + output.getClass().getName());
                    Gson gson = new GsonBuilder().setPrettyPrinting().create();
                    String json = gson.toJson(output);
                    parseAndOutputTranscription(json, model);
                } else {
                    System.out.println("未找到识别结果");
                }
            }
            
        } catch (Exception e) {
            System.err.println("处理识别结果时发生错误: " + e.getMessage());
            e.printStackTrace();
        }
    }
    
    /**
     * 下载识别结果文件
     */
    private static String downloadTranscriptionResult(String transcriptionUrl) {
        try {
            System.out.println("正在下载识别结果...");
            
            HttpURLConnection connection = (HttpURLConnection) new URL(transcriptionUrl).openConnection();
            connection.setRequestMethod("GET");
            connection.setConnectTimeout(10000);
            connection.setReadTimeout(30000);
            connection.connect();
            
            if (connection.getResponseCode() == 200) {
                BufferedReader reader = new BufferedReader(
                    new InputStreamReader(connection.getInputStream(), "UTF-8")
                );
                
                StringBuilder content = new StringBuilder();
                String line;
                while ((line = reader.readLine()) != null) {
                    content.append(line);
                }
                reader.close();
                
                System.out.println("下载完成，文件大小: " + content.length() + " 字符");
                return content.toString();
            } else {
                System.err.println("下载失败，HTTP状态码: " + connection.getResponseCode());
                return null;
            }
            
        } catch (Exception e) {
            System.err.println("下载识别结果时发生错误: " + e.getMessage());
            return null;
        }
    }
    
    /**
     * 解析并输出识别内容
     */
    private static void parseAndOutputTranscription(String jsonContent, String model) {
        try {
            System.out.println("\n解析识别结果 (模型: " + model + ")...");
            
            JsonElement jsonElement = JsonParser.parseString(jsonContent);
            JsonObject jsonObject = jsonElement.getAsJsonObject();
            
            // 尝试不同的结果结构
            if (jsonObject.has("result")) {
                // 结构1: { "result": [...] }
                parseResultStructure(jsonObject, model);
            } else if (jsonObject.has("sentence_list")) {
                // 结构2: { "sentence_list": [...] }
                parseSentenceListStructure(jsonObject, model);
            } else if (jsonObject.has("payload")) {
                // 结构3: { "payload": { ... } }
                parsePayloadStructure(jsonObject, model);
            } else if (jsonObject.has("transcripts")) {
                // 结构4: { "transcripts": [...] }
                parseTranscriptsStructure(jsonObject, model);
            } else {
                // 尝试直接查找sentence_list或其他关键字段
                parseDirectStructure(jsonObject, model);
            }
            
        } catch (Exception e) {
            System.err.println("解析识别内容时发生错误: " + e.getMessage());
            // 安全地截取预览字符串
            int previewLength = Math.min(200, jsonContent.length());
            String preview = previewLength > 0 ? jsonContent.substring(0, previewLength) : "";
            System.err.println("JSON内容预览: " + preview + "...");
            e.printStackTrace();
        }
    }
    
    /**
     * 解析结果结构: { "result": [...] }
     */
    private static void parseResultStructure(JsonObject jsonObject, String model) {
        try {
            JsonArray results = jsonObject.getAsJsonArray("result");
            if (results != null) {
                System.out.println("\n📊 解析到 " + results.size() + " 个结果项");
                
                for (int i = 0; i < results.size(); i++) {
                    JsonObject resultItem = results.get(i).getAsJsonObject();
                    if (resultItem.has("sentence_list")) {
                        parseSentencesWithSpeakers(resultItem.getAsJsonArray("sentence_list"), model);
                        return; // 处理完就返回，避免重复处理
                    } else if (resultItem.has("sentences")) {
                        parseSentencesWithSpeakers(resultItem.getAsJsonArray("sentences"), model);
                        return; // 处理完就返回，避免重复处理
                    } else if (resultItem.has("text")) {
                        System.out.println("文本: " + resultItem.get("text").getAsString());
                    }
                }
            } else {
                System.out.println("结果数组为空");
            }
        } catch (Exception e) {
            System.err.println("解析结果结构时出错: " + e.getMessage());
        }
    }
    
    /**
     * 解析句子列表结构: { "sentence_list": [...] }
     */
    private static void parseSentenceListStructure(JsonObject jsonObject, String model) {
        try {
            if (jsonObject.has("sentence_list")) {
                JsonElement element = jsonObject.get("sentence_list");
                if (element.isJsonArray()) {
                    parseSentencesWithSpeakers(element.getAsJsonArray(), model);
                }
            }
        } catch (Exception e) {
            System.err.println("解析句子列表结构时出错: " + e.getMessage());
        }
    }
    
    /**
     * 解析载荷结构: { "payload": { ... } }
     */
    private static void parsePayloadStructure(JsonObject jsonObject, String model) {
        try {
            JsonObject payload = jsonObject.getAsJsonObject("payload");
            if (payload != null) {
                if (payload.has("sentence_list")) {
                    JsonElement element = payload.get("sentence_list");
                    if (element.isJsonArray()) {
                        parseSentencesWithSpeakers(element.getAsJsonArray(), model);
                    }
                } else if (payload.has("sentences")) {
                    JsonElement element = payload.get("sentences");
                    if (element.isJsonArray()) {
                        parseSentencesWithSpeakers(element.getAsJsonArray(), model);
                    }
                } else if (payload.has("result")) {
                    // 某些情况下可能是payload.result
                    JsonElement resultElement = payload.get("result");
                    if (resultElement.isJsonArray()) {
                        JsonArray results = resultElement.getAsJsonArray();
                        if (results != null && results.size() > 0) {
                            JsonObject firstResult = results.get(0).getAsJsonObject();
                            if (firstResult.has("sentence_list")) {
                                JsonElement element = firstResult.get("sentence_list");
                                if (element.isJsonArray()) {
                                    parseSentencesWithSpeakers(element.getAsJsonArray(), model);
                                }
                            } else if (firstResult.has("sentences")) {
                                JsonElement element = firstResult.get("sentences");
                                if (element.isJsonArray()) {
                                    parseSentencesWithSpeakers(element.getAsJsonArray(), model);
                                }
                            }
                        }
                    }
                }
            }
        } catch (Exception e) {
            System.err.println("解析载荷结构时出错: " + e.getMessage());
        }
    }
    
    /**
     * 解析转录结构: { "transcripts": [...] }
     */
    private static void parseTranscriptsStructure(JsonObject jsonObject, String model) {
        try {
            JsonArray transcripts = jsonObject.getAsJsonArray("transcripts");
            if (transcripts != null && transcripts.size() > 0) {
                for (int i = 0; i < transcripts.size(); i++) {
                    JsonObject transcript = transcripts.get(i).getAsJsonObject();
                    if (transcript.has("sentences")) {
                        JsonElement element = transcript.get("sentences");
                        if (element.isJsonArray()) {
                            parseSentencesWithSpeakers(element.getAsJsonArray(), model);
                            return; // 处理完就返回，避免重复处理
                        }
                    } else if (transcript.has("sentence_list")) {
                        JsonElement element = transcript.get("sentence_list");
                        if (element.isJsonArray()) {
                            parseSentencesWithSpeakers(element.getAsJsonArray(), model);
                            return; // 处理完就返回，避免重复处理
                        }
                    }
                }
            }
        } catch (Exception e) {
            System.err.println("解析转录结构时出错: " + e.getMessage());
        }
    }
    
    /**
     * 直接解析结构
     */
    private static void parseDirectStructure(JsonObject jsonObject, String model) {
        try {
            // 查找所有可能包含句子的字段
            for (String key : jsonObject.keySet()) {
                if (key.contains("sentence") || key.contains("text")) {
                    JsonElement element = jsonObject.get(key);
                    if (element.isJsonArray()) {
                        System.out.println("\n🔍 在字段 '" + key + "' 中发现句子数组");
                        parseSentencesWithSpeakers(element.getAsJsonArray(), model);
                        return;
                    }
                }
            }
            
            // 如果没有找到句子数组，尝试其他方式
            System.out.println("\n📝 未找到标准句子结构，尝试直接输出文本内容:");
            if (jsonObject.has("text")) {
                System.out.println("文本: " + jsonObject.get("text").getAsString());
            } else {
                // 输出所有字符串字段
                for (String key : jsonObject.keySet()) {
                    JsonElement element = jsonObject.get(key);
                    if (element.isJsonPrimitive() && element.getAsJsonPrimitive().isString()) {
                        System.out.println(key + ": " + element.getAsString());
                    }
                }
            }
        } catch (Exception e) {
            System.err.println("直接解析结构时出错: " + e.getMessage());
        }
    }
    
    /**
     * 解析带说话人的句子
     */
    private static void parseSentencesWithSpeakers(JsonArray sentenceList, String model) {
        if (sentenceList == null || sentenceList.size() == 0) {
            System.out.println("未检测到句子");
            return;
        }
        
        System.out.println("\n📝 识别结果（带说话人标记）:");
        System.out.println("模型: " + model);
        System.out.println("==============================");
        
        // 存储说话人信息
        Map<String, String> speakerNames = new HashMap<>();
        int speakerIndex = 1;
        int sentenceCount = 0;
        
        // 遍历句子并输出带说话人标记的内容
        for (int i = 0; i < sentenceList.size(); i++) {
            try {
                JsonElement sentenceElement = sentenceList.get(i);
                if (!sentenceElement.isJsonObject()) {
                    continue;
                }
                
                JsonObject sentence = sentenceElement.getAsJsonObject();
                
                // 获取文本
                String text = "";
                if (sentence.has("text")) {
                    JsonElement textElement = sentence.get("text");
                    if (textElement.isJsonPrimitive()) {
                        text = textElement.getAsString().trim();
                    }
                }
                
                // 跳过空文本
                if (text.isEmpty()) {
                    continue;
                }
                
                // 获取说话人ID
                String speakerId = "unknown";
                if (sentence.has("speaker_id")) {
                    JsonElement speakerElement = sentence.get("speaker_id");
                    if (speakerElement.isJsonPrimitive()) {
                        speakerId = speakerElement.getAsString();
                    }
                } else if (sentence.has("speaker")) {
                    JsonElement speakerElement = sentence.get("speaker");
                    if (speakerElement.isJsonPrimitive()) {
                        speakerId = speakerElement.getAsString();
                    }
                } else if (sentence.has("spk")) {
                    JsonElement speakerElement = sentence.get("spk");
                    if (speakerElement.isJsonPrimitive()) {
                        speakerId = speakerElement.getAsString();
                    }
                }
                
                // 为新说话人分配名称
                if (!speakerNames.containsKey(speakerId)) {
                    speakerNames.put(speakerId, "说话人" + speakerIndex);
                    speakerIndex++;
                }
                
                String speakerName = speakerNames.get(speakerId);
                System.out.println(speakerName + " (ID:" + speakerId + "): " + text);
                sentenceCount++;
            } catch (Exception e) {
                System.err.println("解析第 " + (i+1) + " 个句子时出错: " + e.getMessage());
            }
        }
        
        if (sentenceCount == 0) {
            System.out.println("未检测到有效句子内容");
        }
        
        System.out.println("==============================");
        System.out.println("✅ " + model + " 对话识别完成，共识别出 " + speakerNames.size() + " 位说话人，" + sentenceCount + " 条句子");
        
        // 如果只识别出一个说话人，给出建议
        if (speakerNames.size() <= 1) {
            System.out.println("\n⚠️  注意: " + model + " 只识别出一个说话人");
            System.out.println("   建议尝试以下方法:");
            System.out.println("   1. 调整参数，如降低 min_speaker_duration 值");
            System.out.println("   2. 尝试其他模型");
            System.out.println("   3. 检查音频质量是否足够清晰");
        } else {
            System.out.println("\n🎉 " + model + " 成功识别出多个说话人!");
        }
    }
}
