// script.js

const VIDEO = document.getElementById('webcam');
const CONNECT_SERIAL_BTN = document.getElementById('connectSerialBtn');
const DISCONNECT_SERIAL_BTN = document.getElementById('disconnectSerialBtn');
const LOAD_MODEL_BTN = document.getElementById('loadModelBtn');
const START_WEBCAM_BTN = document.getElementById('startWebcamBtn');
const STOP_WEBCAM_BTN = document.getElementById('stopWebcamBtn');
const MODEL_STATUS = document.getElementById('modelStatus');
const SERIAL_STATUS = document.getElementById('serialStatus');
const PREDICTION_OUTPUT = document.getElementById('prediction');

let mobilenet; // 这个变量将存储加载后的 MobileNet 模型实例
let knnClassifier; // 这个变量将存储 KNN 分类器实例
let classNames = []; // 类别名称将从加载的模型中获取
let webcamStream = null;
let isPredicting = false;

// Web Serial API 变量
let serialPort = null;
let serialWriter = null;
const SERIAL_BAUD_RATE = 9600;
const SERIAL_SEND_MIN_INTERVAL = 500; // 最短发送间隔，500ms，用于内部节流
let lastSerialCommand = ''; // 用于sendToSerialPort内部的节流（避免重复命令+短间隔）
let lastSerialSendTime = 0;

// ===================================
// 新增变量：追踪串口连接状态，并记录上一次发送到串口的类别命令
let isSerialConnectedState = false; // 跟踪串口的逻辑连接状态
let lastSentClassCommand = null;    // 上一次"成功通过确认机制并实际发送"的命令
// ===================================

// ===================================
// 新增变量：用于实现“确认式”发送逻辑
let pendingCommandToSend = null;    // 正在等待确认的命令
let pendingCommandTimerId = null;   // 确认定时器ID
const CONFIRMATION_DELAY_MS = 100; // 等待 100 毫秒确认 class 是否稳定
// ===================================


// ===================================
// Helper Functions (UI Status)
// ===================================
function showStatus(element, type, message) {
    element.className = `status-message status-${type}`;
    element.textContent = message;
}

function updateSerialUI(isConnected) {
    CONNECT_SERIAL_BTN.disabled = isConnected;
    DISCONNECT_SERIAL_BTN.disabled = !isConnected;
    isSerialConnectedState = isConnected; // 更新我们维护的串口连接状态
    if (!isConnected) {
        showStatus(SERIAL_STATUS, 'info', '串口未连接。点击 "连接串口" 开始。');
    }
}

function updateWebcamUI(isRunning) {
    START_WEBCAM_BTN.disabled = isRunning;
    STOP_WEBCAM_BTN.disabled = !isRunning;
}

function updateModelUI(isLoaded) {
    LOAD_MODEL_BTN.disabled = false; // 总是允许重新加载模型
    START_WEBCAM_BTN.disabled = !isLoaded; // 模型加载后才能启动摄像头
}

// ===================================
// Core Logic: Model & Webcam
// ===================================
async function initModel() {
    showStatus(MODEL_STATUS, 'info', '正在加载 MobileNet 模型...');
    try {
        mobilenet = await window.mobilenet.load({ version: 2, alpha: 1.0 });
        knnClassifier = window.knnClassifier.create();
        showStatus(MODEL_STATUS, 'success', 'MobileNet 模型和 KNN 分类器已加载。请加载模型文件。');
        updateModelUI(false); // 此时 MobileNet 已加载，但 KNN 分类器本身还需加载数据
    } catch (error) {
        showStatus(MODEL_STATUS, 'error', `模型加载失败: ${error.message}`);
        console.error('MobileNet/KNN加载失败:', error);
    }
}

// 提取 MobileNet 特征
async function getFeatures(img) {
    if (!mobilenet) {
        throw new Error("MobileNet model is not loaded.");
    }
    return tf.tidy(() => {
        const embeddings = mobilenet.infer(img, true);
        const norm = tf.norm(embeddings);
        const normalized = tf.div(embeddings, norm);
        return normalized;
    });
}

/**
 * 加载 KNN 模型文件，支持 '.json' 和 '.bin' 两部分文件。
 * 用户需要依次选择对应的 .json 和 .bin 文件。
 */
async function loadKNNModel() {
    const inputJson = document.createElement('input');
    inputJson.type = 'file';
    inputJson.accept = '.json';
    inputJson.multiple = false;
 
    showStatus(MODEL_STATUS, 'info', '请先选择 KNN 模型配置文件 (.json)...');
 
    inputJson.onchange = async (e) => {
        const jsonFile = e.target.files[0];
        if (!jsonFile) {
            showStatus(MODEL_STATUS, 'info', '未选择 .json 文件。');
            updateModelUI(false);
            return;
        }
 
        showStatus(MODEL_STATUS, 'info', `正在解析 ${jsonFile.name}...`);
        
        let modelData;
        try {
            const reader = new FileReader();
            const jsonText = await new Promise((resolve, reject) => {
                reader.onload = () => resolve(reader.result);
                reader.onerror = () => reject(reader.error);
                reader.readAsText(jsonFile);
            });
            modelData = JSON.parse(jsonText);
            
            if (!modelData.dataFile) {
                console.warn('模型JSON文件不包含 "dataFile" 字段，尝试以旧的单文件JSON格式加载。');
                return loadSingleJsonModel(modelData);
            }
 
        } catch (error) {
            showStatus(MODEL_STATUS, 'error', `解析 .json 文件失败: ${error.message}`);
            console.error('解析 .json 失败:', error);
            updateModelUI(false);
            return;
        }
 
        const inputBin = document.createElement('input');
        inputBin.type = 'file';
        inputBin.accept = '.bin';
        inputBin.multiple = false;
 
        showStatus(MODEL_STATUS, 'info', `已加载 .json 文件。请选择对应的权重文件 "${modelData.dataFile}" (.bin)...`);
 
        inputBin.onchange = async (eBin) => {
            const binFile = eBin.target.files[0];
            if (!binFile) {
                showStatus(MODEL_STATUS, 'info', '未选择 .bin 文件。');
                updateModelUI(false);
                return;
            }
 
            if (binFile.name !== modelData.dataFile) {
                showStatus(MODEL_STATUS, 'error', `选择的 .bin 文件名 "${binFile.name}" 与 .json 中定义的 "${modelData.dataFile}" 不匹配！请选择正确的文件。`);
                updateModelUI(false);
                return;
            }
 
            showStatus(MODEL_STATUS, 'info', `正在读取 ${binFile.name} (二进制权重文件)...`);
            let binData;
            try {
                const reader = new FileReader();
                const arrayBuffer = await new Promise((resolve, reject) => {
                    reader.onload = () => resolve(reader.result);
                    reader.onerror = () => reject(reader.error);
                    reader.readAsArrayBuffer(binFile);
                });
                binData = new Float32Array(arrayBuffer);
            } catch (error) {
                showStatus(MODEL_STATUS, 'error', `读取 .bin 文件失败: ${error.message}`);
                console.error('读取 .bin 失败:', error);
                updateModelUI(false);
                return;
            }
 
            try {
                knnClassifier.clearAllClasses();
 
                Object.keys(modelData.dataset).forEach(label => {
                    const classDataMeta = modelData.dataset[label];
                    const startFloat32ElementIndex = classDataMeta.start;
                    const numFloat32Elements = classDataMeta.length;
                    
                    const featureDim = modelData.featureDim || 1280;
 
                    const classFeatures = binData.subarray(startFloat32ElementIndex, startFloat32ElementIndex + numFloat32Elements);
 
                    if (classFeatures.length === 0) {
                        console.warn(`类别 ${label} 没有找到特征数据，跳过。`);
                        return;
                    }
 
                    if (classFeatures.length % featureDim !== 0) {
                        const actualSamples = classFeatures.length / featureDim;
                        console.error(
                            `--- 类别: ${label} ---`,
                            `起始 Float32 元素索引: ${startFloat32ElementIndex}`,
                            `该类别 Float32 元素数量: ${numFloat32Elements}`,
                            `ERROR: 特征数据长度 (${classFeatures.length} 个 Float32 元素) 与特征维度 (${featureDim}) 不匹配！` +
                            `实际样本数计算为 ${actualSamples} (预期为整数)。`,
                            `请检查您的模型导出逻辑和训练数据的完整性。`
                        );
                        throw new Error("模型数据完整性错误：特征数据长度与维度不匹配。");
                    }
                    
                    const numSamples = classFeatures.length / featureDim;
 
                    for (let i = 0; i < numSamples; i++) {
                        const startIndex = i * featureDim;
                        const endIndex = (i + 1) * featureDim;
                        const sampleFeatures = classFeatures.subarray(startIndex, endIndex);
 
                        const sampleTensor = tf.tensor(sampleFeatures, [1, featureDim]);
                        
                        knnClassifier.addExample(sampleTensor, parseInt(label));
                        tf.dispose(sampleTensor);
                    }
                });
 
                if (modelData.classList && Array.isArray(modelData.classList)) {
                    classNames = modelData.classList.map(c => c.name);
                } else {
                    console.warn('模型JSON中未找到 classList 字段或格式不正确，使用默认类别名称。');
                    classNames = Object.keys(modelData.dataset).map(key => `Class ${parseInt(key) + 1}`);
                }
 
                showStatus(MODEL_STATUS, 'success', `模型 "${jsonFile.name}" 及权重加载成功！类别: ${classNames.join(', ')}。`);
                updateModelUI(true);
                
            } catch (error) {
                showStatus(MODEL_STATUS, 'error', `处理模型数据失败: ${error.message}`);
                console.error('处理模型数据失败:', error);
                updateModelUI(false);
            }
        };
        inputBin.click();
    };
    inputJson.click();
}
 
/**
 * 辅助函数：处理旧的单文件JSON模型格式（ dataset 字段直接包含数据而不是偏移量）
 * @param {object} modelData - 已解析的 JSON 模型数据
 */
async function loadSingleJsonModel(modelData) {
    try {
        knnClassifier.clearAllClasses();
        Object.keys(modelData.dataset).forEach(key => {
            const data = modelData.dataset[key];
            const featureDim = modelData.featureDim || 1280;
            if (data.length % featureDim !== 0) {
                throw new Error(`类别 ${key} 的特征数据长度 ${data.length} 与特征维度 ${featureDim} 不匹配！`);
            }
            const numSamples = data.length / featureDim;
            const tensor = tf.tensor(data, [numSamples, featureDim]);
            knnClassifier.addExample(tensor, parseInt(key));
        });
 
        if (modelData.classList && Array.isArray(modelData.classList)) {
            classNames = modelData.classList.map(c => c.name);
        } else if (modelData.classNames && Array.isArray(modelData.classNames)) {
            classNames = modelData.classNames;
        } else {
            console.warn('模型JSON中未找到 classList/classNames 字段，使用默认类别名称。');
            classNames = Object.keys(modelData.dataset).map(key => `Class ${parseInt(key) + 1}`);
        }
        
        showStatus(MODEL_STATUS, 'success', `模型 (单文件JSON格式) 加载成功！类别: ${classNames.join(', ')}。`);
        updateModelUI(true);
    } catch (error) {
        showStatus(MODEL_STATUS, 'error', `加载单文件JSON模型失败: ${error.message}`);
        console.error('加载单文件JSON模型失败:', error);
        updateModelUI(false);
    }
}


async function startWebcam() {
    if (webcamStream) return;

    if (!knnClassifier || knnClassifier.getNumClasses() === 0) {
        showStatus(MODEL_STATUS, 'error', '请先加载训练好的模型！');
        return;
    }

    try {
        const stream = await navigator.mediaDevices.getUserMedia({ video: { facingMode: 'user' }, audio: false });
        VIDEO.srcObject = stream;
        webcamStream = stream;
        updateWebcamUI(true);

        VIDEO.onloadeddata = () => {
             // ===================================
            // 启动摄像头时，重置所有发送状态，确保第一次识别结果可发送
            if (pendingCommandTimerId) {
                clearTimeout(pendingCommandTimerId);
                pendingCommandTimerId = null;
            }
            pendingCommandToSend = null;
            lastSentClassCommand = null; 
            // ===================================
            isPredicting = true;
            predictLoop();
            showStatus(MODEL_STATUS, 'info', '摄像头已启动，开始实时预测...');
        };

    } catch (error) {
        showStatus(MODEL_STATUS, 'error', `无法访问摄像头: ${error.message}`);
        console.error('启动摄像头失败:', error);
        updateWebcamUI(false);
    }
}

function stopWebcam() {
    if (webcamStream) {
        webcamStream.getTracks().forEach(track => track.stop());
        webcamStream = null;
    }
    isPredicting = false;
    VIDEO.srcObject = null;
    updateWebcamUI(false);
    PREDICTION_OUTPUT.textContent = '等待识别...';
    showStatus(MODEL_STATUS, 'info', '摄像头已停止。');

    // ===================================
    // 停止摄像头时，清除任何待确认的命令，并发送“停止”或“复位”命令
    if (pendingCommandTimerId) {
        clearTimeout(pendingCommandTimerId);
        pendingCommandTimerId = null;
    }
    pendingCommandToSend = null;

    // 假设 '0' 是停止/复位命令，并且只有当上次发送的不是 '0' 时才发送
    if (lastSentClassCommand !== '0') { 
        sendToSerialPort('0'); 
    }
    lastSentClassCommand = null; // 重置，确保下次启动摄像头时能发送初始状态
    // ===================================
}

// ============== 重要的修改区域 start ==============
// 定义一个全局变量，用于存储当前帧识别到的命令，供 setTimeout 回调使用
let currentCommandInFrame = '0'; // 初始值默认是 '0' (复位/无动作)

async function predictLoop() {
    if (!isPredicting) return;

    if (VIDEO.readyState === 4 && VIDEO.videoWidth > 0 && VIDEO.videoHeight > 0) {
        let commandCandidate = '0'; // 每次循环开始时，将当前帧的候选命令初始化为'0'

        try {
            const features = await getFeatures(VIDEO);
            const k = 3;

            if (!knnClassifier || knnClassifier.getNumClasses() === 0) {
                features.dispose();
                PREDICTION_OUTPUT.textContent = 'KNN 分类器未就绪或无数据。';
                commandCandidate = '0'; // 使用默认命令
            } else {
                const prediction = await knnClassifier.predictClass(features, k);
                features.dispose();

                if (prediction && prediction.confidences) {
                    let maxConfidence = 0;
                    let predictedClassIndex = -1;
                    
                    const confidencesArray = Object.entries(prediction.confidences).map(([key, value]) => ({ index: parseInt(key), confidence: value }));
                    
                    confidencesArray.forEach(({ index, confidence }) => {
                        if (confidence > maxConfidence) {
                            maxConfidence = confidence;
                            predictedClassIndex = index;
                        }
                    });

                    const confidenceThreshold = 0.75; // 75%置信度
                    if (predictedClassIndex !== -1 && maxConfidence > confidenceThreshold) {
                        const className = classNames[predictedClassIndex] || `Class ${predictedClassIndex + 1}`;
                        const percentage = (maxConfidence * 100).toFixed(1);
                        PREDICTION_OUTPUT.textContent = `识别为: ${className} (${percentage}%)`;

                        // 根据类别设置本帧的候选命令
                        if (predictedClassIndex === 0) {
                            commandCandidate = '1';
                        } else if (predictedClassIndex === 1) {
                            commandCandidate = '2';
                        } else {
                            commandCandidate = '0'; // 未匹配到特定类别，或默认复位
                        }
                    } else {
                        PREDICTION_OUTPUT.textContent = `未知或不确定... (最高置信度: ${(maxConfidence * 100).toFixed(1)}%)`;
                        commandCandidate = '0'; // 不确定也发送'0'回退
                    }
                } else {
                    PREDICTION_OUTPUT.textContent = '无法识别。';
                    commandCandidate = '0'; // 无法识别也发送 '0' 回退
                }
            }
        } catch (error) {
            console.error('预测错误:', error);
            PREDICTION_OUTPUT.textContent = `预测错误: ${error.message}`;
            commandCandidate = '0'; // 错误时也发送'0'
        }

        // =========================================================
        // 核心逻辑：确认式发送机制 (Persistence / Confirmation)
        // =========================================================

        // 更新全局变量 currentCommandInFrame，指示当前帧的预测结果
        currentCommandInFrame = commandCandidate;

        // 只有当当前帧的候选命令与上一次实际发送的命令不同时，才进入确认流程
        if (currentCommandInFrame !== lastSentClassCommand) {
            // 如果目前没有待确认的命令，或者待确认的命令与当前不同
            if (pendingCommandToSend === null || pendingCommandToSend !== currentCommandInFrame) {
                if (pendingCommandTimerId) {
                    clearTimeout(pendingCommandTimerId); // 清除旧的等待定时器
                    pendingCommandTimerId = null;
                    // console.log(`[${new Date().toLocaleTimeString()}] 命令在确认前再次变化，取消旧的确认定时器.`);
                }
                
                pendingCommandToSend = currentCommandInFrame; // 设置新的待确认命令
                
                // 设置一个定时器，在 CONFIRMATION_DELAY_MS 后进行确认
                pendingCommandTimerId = setTimeout(async () => {
                    // 当定时器触发时，再次检查全局的 currentCommandInFrame
                    // 如果它仍然与我们开始等待的命令相同，则确认并发送
                    if (currentCommandInFrame === pendingCommandToSend) {
                        console.log(`[${new Date().toLocaleTimeString()}] 确认命令 "${pendingCommandToSend}" 稳定，正在发送。`);
                        await sendToSerialPort(pendingCommandToSend);
                        lastSentClassCommand = pendingCommandToSend; // 更新上一次实际发送的命令
                    } else {
                        console.log(`[${new Date().toLocaleTimeString()}] 命令在确认期内再次变化 (${pendingCommandToSend} -> ${currentCommandInFrame})，不发送。`);
                    }
                    pendingCommandToSend = null;    // 清理待确认状态
                    pendingCommandTimerId = null;   // 清理定时器ID
                }, CONFIRMATION_DELAY_MS);

                // console.log(`[${new Date().toLocaleTimeString()}] 检测到新候选命令 "${currentCommandInFrame}"，启动 ${CONFIRMATION_DELAY_MS}ms 确认计时器。`);
            }
            // 如果 currentCommandInFrame 1. 仍然与 pendingCommandToSend 相同
            //    并且 2. 不同于 lastSentClassCommand
            // 说明它处在“正在等待确认”的状态，不需要做任何事情，让当前定时器继续运行
        } else { // 当前帧的命令与上次实际发送的命令相同
            // 如果已经发送了相同的命令，或者回到了已发送的命令，那么清除待确认定时器
            if (pendingCommandTimerId) {
                clearTimeout(pendingCommandTimerId);
                pendingCommandTimerId = null;
                pendingCommandToSend = null; // 清理待确认状态
                // console.log(`[${new Date().toLocaleTimeString()}] 当前命令 "${currentCommandInFrame}" 与上次发送命令相同，或回摆，取消所有待确认发送。`);
            }
        }
        // =========================================================
    }
    requestAnimationFrame(predictLoop); // 继续下一帧预测
}
// ============== 重要的修改区域 end ==============


// ===================================
// Web Serial API Logic
// ===================================
async function checkWebSerialCompatibility() {
    if ('serial' in navigator) {
        showStatus(SERIAL_STATUS, 'info', 'Web Serial API 可用！');
        CONNECT_SERIAL_BTN.disabled = false;
    } else {
        showStatus(SERIAL_STATUS, 'error', 'Web Serial API 在此浏览器中不可用。请使用最新版 Chrome！');
        CONNECT_SERIAL_BTN.disabled = true;
    }
}

async function connectSerial() {
    showStatus(SERIAL_STATUS, 'info', '正在请求连接串口...');
    try {
        serialPort = await navigator.serial.requestPort();
        
        serialPort.addEventListener('disconnect', () => {
             const disconnectTime = new Date().toLocaleString();
             console.warn(`[${disconnectTime}] ❗️❗️❗️ 串口已断开连接事件触发！ `); // 强调日志
            showStatus(SERIAL_STATUS, 'error', '串口连接已丢失！');
            disconnectSerial();
        });
 
        await serialPort.open({ baudRate: SERIAL_BAUD_RATE });
        console.log('串口已成功打开。尝试获取写入器...');
        serialWriter = serialPort.writable.getWriter();
        console.log('写入器已获取。');
        
        console.log('串口已连接并打开。');
        showStatus(SERIAL_STATUS, 'success', `串口已连接 (Baud: ${SERIAL_BAUD_RATE})`);
        updateSerialUI(true); // 更新状态为已连接
    } catch (error) {
        console.error('连接串口失败:', error);
        if (error.name === 'NotFoundError') {
            showStatus(SERIAL_STATUS, 'warning', '连接串口请求已取消或未选择端口。请选择一个设备。');
        } else if (error.name === 'SecurityError') {
            showStatus(SERIAL_STATUS, 'error', `连接串口失败: ${error.message}。请确保您在安全上下文 (HTTPS 或 localhost) 中运行。`);
        } else {
            showStatus(SERIAL_STATUS, 'error', `连接串口失败: ${error.message}`);
        }
        updateSerialUI(false); // 更新状态为未连接
    }
}

async function disconnectSerial() {
    console.trace(`[${new Date().toLocaleString()}] disconnectSerial() called`);
        
    if (!isSerialConnectedState && !serialPort && !serialWriter) {
        console.log('Already in disconnected state, skipping further cleanup.');
        return;
    }

    // ===================================
    // 断开连接时，清除任何待确认的定时器和状态
    if (pendingCommandTimerId) {
        clearTimeout(pendingCommandTimerId);
        pendingCommandTimerId = null;
    }
    pendingCommandToSend = null;
    // ===================================

    try {
        if (serialWriter) {
            await serialWriter.close();
            serialWriter = null;
        }
        if (serialPort && serialPort.readable) {
            const reader = serialPort.readable.getReader();
            if (reader) {
                await reader.cancel();
                reader.releaseLock();
            }
        }
        if (serialPort && serialPort.isOpen) {
            await serialPort.close();
        }
        serialPort = null;
        console.log('串口已断开。');
        showStatus(SERIAL_STATUS, 'info', '串口已断开。');
    } catch (error) {
        console.error('断开串口失败或串口已处于非连接状态:', error);
        showStatus(SERIAL_STATUS, 'error', `断开串口失败: ${error.message}`);
    } finally {
        updateSerialUI(false); // 无论如何都确保UI和状态更新为断开
        lastSentClassCommand = null; // 断开连接时，重置上一次发送的命令，确保重连后第一次变化能发送
    }
}

/**
 * 低级串口发送函数。包含内部的频繁发送节流。
 * @param {string} command - 要发送的命令字符串。
 */
async function sendToSerialPort(command) {
    if (!isSerialConnectedState) {
        // console.warn('串口已断开（逻辑状态），无法发送命令。'); 
        return;
    }

    if (!serialWriter) {
        console.warn(`[${new Date().toLocaleString()}] 串口连接状态为 true，但 serialWriter 不可用。尝试重置串口连接。`);
        disconnectSerial(); 
        return;
    }

    // 避免频繁发送相同的命令 (sendToSerialPort 内部的节流)
    const currentTime = new Date().getTime();
    if (command === lastSerialCommand && (currentTime - lastSerialSendTime < SERIAL_SEND_MIN_INTERVAL)) {
        return; // 命令相同且发送间隔太短，不发送
    }

    try {
        const encoder = new TextEncoder();
        const data = encoder.encode(command + '\n');
        await serialWriter.write(data);
        console.log(`[${new Date().toLocaleTimeString()}] 实际发送串口命令: ${command}`); 
        lastSerialCommand = command; // 更新内部节流用的命令
        lastSerialSendTime = currentTime;
    } catch (error) {
        console.error(`[${new Date().toLocaleTimeString()}] 实际发送串口命令失败: ${error.message}`);
        showStatus(SERIAL_STATUS, 'error', `发送串口命令失败: ${error.message}，正在断开串口。`);
        disconnectSerial(); 
    }
}

// ===================================
// Event Listeners
// ===================================
CONNECT_SERIAL_BTN.addEventListener('click', connectSerial);
DISCONNECT_SERIAL_BTN.addEventListener('click', disconnectSerial);
LOAD_MODEL_BTN.addEventListener('click', loadKNNModel);
START_WEBCAM_BTN.addEventListener('click', startWebcam);
STOP_WEBCAM_BTN.addEventListener('click', stopWebcam);

// ===================================
// Initialization
// ===================================
document.addEventListener('DOMContentLoaded', () => {
    checkWebSerialCompatibility(); // 检查 Web Serial API 兼容性
    initModel(); // 加载 MobileNet 和 KNN 分类器实例
});