// 全局变量和模型实例
let recognizer; // 基础的 SpeechCommands recognizer
let transferRecognizer; // 用于迁移学习的 recognizer
const labels = []; // 用户定义的类别标签数组 (包括背景噪音)
// 将背景噪音定义为第一个类别，其内部名称为 _background_noise_
const BACKGROUND_NOISE_LABEL = '_background_noise_';
const BACKGROUND_NOISE_INDEX = 0; // 仅用于本地 labels 数组索引，不直接用于collectExample

let isPredicting = false; // 预测状态标志
let isRecording = false; // 录音状态标志，防止重复点击
const recordDuration = 1000; // 每个样本的录音时长 (毫秒)
let isModelTrainedFlag = false; // 手动维护模型训练状态
let predictionStopFunction = null; // 存储 transferRecognizer.listen() 返回的停止函数

/**
 * 初始化函数 - 加载模型和创建迁移学习模型
 * @returns {Promise<void>}
 */
async function init() {
    try {
        recognizer = speechCommands.create(
            'BROWSER_FFT' // 使用浏览器内置的 FFT 处理，性能更好
        );

        await recognizer.ensureModelLoaded();
        transferRecognizer = recognizer.createTransfer('my-custom-model');
        
        // 只有在 transferRecognizer 创建成功后，才将背景噪音标签加入我们的 local labels 数组
        labels.push(BACKGROUND_NOISE_LABEL);
        
        return Promise.resolve();
    } catch (error) {
        return Promise.reject(error);
    }
}

/**
 * 批量录制样本的通用函数
 * @param {string} label - 标签名称
 * @param {number} countToRecord - 要录制的样本数量
 * @returns {Promise<void>}
 */
async function recordMultipleExamples(label, countToRecord = 5) {
    if (isRecording) {
        return Promise.reject(new Error('正在录制中，请等待当前录音完成'));
    }

    isRecording = true;
    
    for (let i = 0; i < countToRecord; i++) {
        try {
            await transferRecognizer.collectExample(
                label,
                { amplitudeRequired: true, durationMillis: recordDuration }
            );
            // 在每次录音之间增加短暂延迟，以便更好地分离样本
            if (i < countToRecord - 1) {
                await new Promise(resolve => setTimeout(resolve, Math.max(200, recordDuration / 5)));
            }
        } catch (error) {
            isRecording = false;
            return Promise.reject(error);
        }
    }

    isRecording = false;
    return Promise.resolve();
}

/**
 * 添加自定义类别
 * @param {string} categoryName - 类别名称
 */
function addCustomCategory(categoryName) {
    if (!categoryName) {
        return Promise.reject(new Error('类别名称不能为空'));
    }
    
    // 检查是否与现有标签重复
    if (labels.some(label => label.toLowerCase() === categoryName.toLowerCase())) {
        return Promise.reject(new Error(`类别 "${categoryName}" 已经存在`));
    }

    // 将标签添加到本地数组
    labels.push(categoryName);
    return Promise.resolve();
}

/**
 * 检查训练就绪状态
 * @returns {boolean} 是否可以开始训练
 */
function checkTrainingReadiness() {
    const exampleCounts = transferRecognizer.countExamples();
    
    let backgroundNoiseReady = (exampleCounts[BACKGROUND_NOISE_LABEL] || 0) > 0;
    
    let customCategoriesReady = 0;
    // 遍历本地 labels 数组，检查每个自定义类别是否有样本
    for (let i = 1; i < labels.length; i++) { // 从索引 1 开始，因为 0 是背景噪音
        const customLabel = labels[i];
        if ((exampleCounts[customLabel] || 0) > 0) {
            customCategoriesReady++;
        }
    }

    // 必须有背景噪音样本，并且至少一个自定义类别有样本
    return backgroundNoiseReady && customCategoriesReady >= 1;
}

/**
 * 模型训练函数
 * @param {Object} trainingConfig - 训练配置参数
 * @returns {Promise<void>}
 */
async function trainModel(trainingConfig = {}) {
    const exampleCounts = transferRecognizer.countExamples();
    
    let totalExamples = 0;
    let validClasses = 0;
    const MIN_SAMPLES_PER_CLASS_FOR_TRAINING = 5; 
    let allClassesHaveEnoughSamples = true; 
    
    // 统计所有类别的有效样本数，并检查每个类别是否达到最低要求
    for (const labelName of labels) {
        if (exampleCounts[labelName] && exampleCounts[labelName] > 0) {
            totalExamples += exampleCounts[labelName];
            validClasses++;
            if (exampleCounts[labelName] < MIN_SAMPLES_PER_CLASS_FOR_TRAINING) {
                allClassesHaveEnoughSamples = false; 
            }
        }
    }

    if (validClasses < 2) { 
        return Promise.reject(new Error(`训练需要至少 "背景噪音" 和另一个自定义类别。当前只有 ${validClasses} 个有效类别。`));
    }
    
    if (!allClassesHaveEnoughSamples) {
        return Promise.reject(new Error(`请确保每个类别至少收集了 ${MIN_SAMPLES_PER_CLASS_FOR_TRAINING} 个样本。`));
    }
    
    if (totalExamples === 0) {
        return Promise.reject(new Error('没有收集到任何训练样本'));
    }

    const defaultConfig = {
        epochs: 50, 
        batchSize: 16, 
        validationSplit: 0.1, 
        shuffle: true, 
        yieldEvery: 'epoch'
    };
    
    const config = Object.assign(defaultConfig, trainingConfig);

    try {
        await transferRecognizer.train(config);
        isModelTrainedFlag = true;
        return Promise.resolve();
    } catch (error) {
        isModelTrainedFlag = false;
        return Promise.reject(error);
    }
}

/**
 * 开始实时预测
 * @param {Function} onPrediction - 预测结果回调函数
 * @param {Object} listenOptions - 监听选项
 * @returns {Promise<Function>} 停止预测的函数
 */
async function startPrediction(onPrediction, listenOptions = {}) {
    if (isPredicting) {
        return Promise.reject(new Error('识别已经在进行中'));
    }
    
    if (!isModelTrainedFlag) { 
        return Promise.reject(new Error('模型尚未训练完成'));
    }

    isPredicting = true;
    
    const defaultOptions = {
        includeEmbedding: true, 
        probabilityThreshold: 0.75, 
        suppressionTimeMillis: 300, 
        overlapFactor: 0.50, 
    };
    
    const options = Object.assign(defaultOptions, listenOptions);

    predictionStopFunction = await transferRecognizer.listen(result => {
        if (!isPredicting) return;

        const classLabels = transferRecognizer.wordLabels(); 
        const scores = result.scores; 
        const maxScore = Math.max(...scores);
        const predictedIndex = scores.indexOf(maxScore);
        
        let predictedLabel = classLabels[predictedIndex];
        // 如果预测结果是内部的背景噪音标签，转换成用户友好的显示
        if (predictedLabel === BACKGROUND_NOISE_LABEL) {
            predictedLabel = '背景噪音'; 
        }
        
        // 调用回调函数返回预测结果
        if (typeof onPrediction === 'function') {
            onPrediction({
                label: predictedLabel,
                score: maxScore,
                scores: scores,
                labels: classLabels.map(label => label === BACKGROUND_NOISE_LABEL ? '背景噪音' : label)
            });
        }

    }, options);

    return Promise.resolve(predictionStopFunction);
}

/**
 * 停止实时预测
 */
function stopPrediction() {
    if (isPredicting) {
        if (typeof predictionStopFunction === 'function') {
            predictionStopFunction();
            predictionStopFunction = null;
        }
        
        isPredicting = false;
    }
}

/**
 * 获取各类别样本数量
 * @returns {Object} 各类别样本数量统计
 */
function getExampleCounts() {
    return transferRecognizer.countExamples();
}

/**
 * 获取模型是否已训练的状态
 * @returns {boolean} 模型是否已训练
 */
function isModelTrained() {
    return isModelTrainedFlag;
}

// 导出公共接口
window.AudioClassifier = {
    init,
    recordMultipleExamples,
    addCustomCategory,
    checkTrainingReadiness,
    trainModel,
    startPrediction,
    stopPrediction,
    getExampleCounts,
    isModelTrained,
    labels
};