// server/utils/textRecognizer.js
import { Tensor } from 'onnxruntime-node';
import sharp from 'sharp';
import fse from 'fs-extra';
import * as path from 'path';

class TextRecognizer {
    constructor() {
        this.recSession = null;
        this.config = null;
        this.characterSet = [];
        this.debugDir = path.join(process.cwd(), 'temp', 'debug');
        this.preprocessedDir = path.join(process.cwd(), 'temp', 'preprocessed');
        this.logger = {
            info: (msg, ...args) => console.log(`🔤 [识别] ${msg}`, ...args),
            error: (msg, ...args) => console.error(`❌ [识别] ${msg}`, ...args),
            debug: (msg, ...args) => console.log(`🐛 [识别] ${msg}`, ...args),
            warn: (msg, ...args) => console.warn(`🐛 [识别] ${msg}`, ...args)
        };

        // 确保目录存在
        fse.ensureDirSync(this.debugDir);
        fse.ensureDirSync(this.preprocessedDir);
    }

    initialize(recSession, config) {
        this.recSession = recSession;
        this.config = config;
        this.logger.info('文本识别器初始化完成');
    }

    async loadCharacterSet(keysPath) {
        try {
            const keysContent = await fse.readFile(keysPath, 'utf8');
            this.characterSet = [];
            const lines = keysContent.split('\n');

            // 使用提供的字符集文件
            const uniqueChars = new Set();

            for (const line of lines) {
                const trimmed = line.trim();
                // 跳过空行和注释行
                if (trimmed && !trimmed.startsWith('#')) {
                    // 将每行作为一个完整的字符处理
                    uniqueChars.add(trimmed);
                }
            }

            this.characterSet = Array.from(uniqueChars);

            if (this.characterSet.length === 0) {
                throw new Error('字符集文件为空或格式不正确');
            }

            this.logger.info(`字符集加载完成: ${this.characterSet.length}个字符`);

            // 记录字符集统计信息
            const charTypes = {
                chinese: 0,
                english: 0,
                digit: 0,
                punctuation: 0,
                other: 0
            };

            this.characterSet.forEach(char => {
                if (/[\u4e00-\u9fff]/.test(char)) {
                    charTypes.chinese++;
                } else if (/[a-zA-Z]/.test(char)) {
                    charTypes.english++;
                } else if (/[0-9]/.test(char)) {
                    charTypes.digit++;
                } else if (/[，。！？；：""（）【】《》…—·]/.test(char)) {
                    charTypes.punctuation++;
                } else {
                    charTypes.other++;
                }
            });

            this.logger.debug(`字符集统计: 中文${charTypes.chinese}, 英文${charTypes.english}, 数字${charTypes.digit}, 标点${charTypes.punctuation}, 其他${charTypes.other}`);
            this.logger.debug(`前20个字符: ${this.characterSet.slice(0, 20).join('')}`);

        } catch (error) {
            this.logger.error('加载字符集失败', error.message);
            // 完全使用提供的字符集，失败时抛出错误
            throw new Error(`字符集加载失败: ${error.message}`);
        }
    }
    getCharacterSetSize() {
        return this.characterSet.length;
    }

    async recognizeText(textRegionBuffer, regionIndex = 0) {
        const startTime = Date.now();
        this.logger.info(`开始文本识别 - 区域 ${regionIndex}`);

        try {
            const inputTensor = await this.prepareRecognitionInput(textRegionBuffer, regionIndex);
            const outputs = await this.recSession.run({ [this.recSession.inputNames[0]]: inputTensor });
            const result = this.postprocessRecognition(outputs);

            const processingTime = Date.now() - startTime;
            this.logger.info(`识别完成 - 区域 ${regionIndex}: "${result.text}", 置信度: ${result.confidence.toFixed(4)}, 耗时: ${processingTime}ms`);

            return result;

        } catch (error) {
            this.logger.error(`文本识别失败 - 区域 ${regionIndex}`, error);
            return { text: '', confidence: 0 };
        }
    }

    async prepareRecognitionInput(textRegionBuffer, regionIndex = 0) {
        this.logger.debug(`准备识别输入 - 区域 ${regionIndex}`);

        const targetHeight = 48;
        const targetWidth = 320; // 原始目标宽度
        const finalWidth = targetWidth + 20; // 最终宽度（左右各加10像素）
        const timestamp = Date.now();

        try {
            const metadata = await sharp(textRegionBuffer).metadata();
            this.logger.debug(`原始区域 ${regionIndex}: ${metadata.width}x${metadata.height}`);

            // 保存原始裁剪区域图像
            const originalPath = path.join(this.preprocessedDir, `region-${regionIndex}-original-${timestamp}.png`);
            await fse.writeFile(originalPath, textRegionBuffer);
            this.logger.debug(`保存原始区域图像: ${originalPath}`);

            // 图像分析
            const stats = await sharp(textRegionBuffer).grayscale().stats();
            const meanBrightness = stats.channels[0].mean;
            const stdDev = stats.channels[0].stdev;

            this.logger.debug(`图像统计 - 区域 ${regionIndex}: 亮度=${meanBrightness.toFixed(1)}, 对比度=${stdDev.toFixed(1)}`);

            // 智能预处理
            let processedBuffer = await this.applySmartPreprocessing(textRegionBuffer, meanBrightness, stdDev, regionIndex);

            // 保存预处理后的图像（灰度+对比度调整后）
            const processedPath = path.join(this.preprocessedDir, `region-${regionIndex}-processed-${timestamp}.png`);
            await fse.writeFile(processedPath, processedBuffer);
            this.logger.debug(`保存预处理图像: ${processedPath}`);

            // 保持宽高比的resize，并在左右添加10像素空白
            const resizedBuffer = await this.resizeWithAspectRatio(processedBuffer, targetWidth, targetHeight, regionIndex);

            // 保存调整大小后的图像
            const resizedPath = path.join(this.preprocessedDir, `region-${regionIndex}-resized-${timestamp}.png`);
            await fse.writeFile(resizedPath, resizedBuffer);
            this.logger.debug(`保存调整大小图像: ${resizedPath}`);

            // 使用最终尺寸创建张量
            const inputData = await this.bufferToTensor(resizedBuffer, finalWidth, targetHeight);
            this.logger.debug(`识别输入张量准备完成 - 区域 ${regionIndex}`);

            // 创建张量时使用最终尺寸
            return new Tensor('float32', inputData, [1, 3, targetHeight, finalWidth]);

        } catch (error) {
            this.logger.error(`准备识别输入失败 - 区域 ${regionIndex}`, error);
            return new Tensor('float32', new Float32Array(3 * targetHeight * finalWidth).fill(0.5), [1, 3, targetHeight, finalWidth]);
        }
    }

// server/utils/textRecognizer.js
// 增强的图像预处理

    async applySmartPreprocessing(buffer, meanBrightness, stdDev, regionIndex = 0) {
        let processedBuffer = buffer;

        try {
            // 更精细的图像分析
            const stats = await sharp(buffer)
                .grayscale()
                .stats();

            const median = stats.channels[0].median;
            const max = stats.channels[0].max;
            const min = stats.channels[0].min;

            this.logger.debug(`区域 ${regionIndex}: 详细统计 - 中值=${median}, 范围=${min}-${max}, 均值=${meanBrightness.toFixed(1)}, 标准差=${stdDev.toFixed(1)}`);

            // 更智能的预处理策略
            if (meanBrightness > 220 && stdDev < 25) {
                // 高亮度低对比度图像
                this.logger.debug(`区域 ${regionIndex}: 应用高亮度低对比度增强`);
                processedBuffer = await sharp(buffer)
                    .linear(1.8, -80)  // 更强的对比度增强
                    .normalize({ lower: 5, upper: 95 }) // 更激进的归一化
                    .grayscale()
                    .toBuffer();
            } else if (meanBrightness < 70) {
                // 低亮度图像
                this.logger.debug(`区域 ${regionIndex}: 应用低亮度增强`);
                processedBuffer = await sharp(buffer)
                    .linear(1.5, 50)   // 更强的亮度提升
                    .normalize()
                    .grayscale()
                    .toBuffer();
            } else if (stdDev < 15) {
                // 极低对比度
                this.logger.debug(`区域 ${regionIndex}: 应用极低对比度增强`);
                processedBuffer = await sharp(buffer)
                    .linear(2.0, -30)  // 非常强的对比度增强
                    .normalize({ lower: 1, upper: 99 })
                    .grayscale()
                    .toBuffer();
            } else if (stdDev > 80) {
                // 高对比度图像，可能过度增强
                this.logger.debug(`区域 ${regionIndex}: 应用高对比度抑制`);
                processedBuffer = await sharp(buffer)
                    .linear(0.8, 20)   // 降低对比度
                    .normalize()
                    .grayscale()
                    .toBuffer();
            } else {
                // 标准处理
                this.logger.debug(`区域 ${regionIndex}: 应用标准增强`);
                processedBuffer = await sharp(buffer)
                    .linear(1.3, -15)  // 适度的对比度增强
                    .normalize({ lower: 10, upper: 90 })
                    .grayscale()
                    .toBuffer();
            }

            // 应用锐化滤波增强文字边缘
            processedBuffer = await sharp(processedBuffer)
                .sharpen({
                    sigma: 1.2,
                    m1: 1.5,
                    m2: 0.7
                })
                .toBuffer();

        } catch (error) {
            this.logger.error(`区域 ${regionIndex}: 预处理失败`, error);
            // 回退到基本处理
            processedBuffer = await sharp(buffer)
                .normalize()
                .grayscale()
                .toBuffer();
        }

        return processedBuffer;
    }

    async resizeWithAspectRatio(buffer, targetWidth, targetHeight, regionIndex = 0) {
        const metadata = await sharp(buffer).metadata();
        const originalAspectRatio = metadata.width / metadata.height;
        const targetAspectRatio = targetWidth / targetHeight;

        let resizeWidth, resizeHeight;

        if (originalAspectRatio > targetAspectRatio) {
            // 宽度限制，按宽度缩放
            resizeWidth = targetWidth;
            resizeHeight = Math.round(targetWidth / originalAspectRatio);
        } else {
            // 高度限制，按高度缩放
            resizeHeight = targetHeight;
            resizeWidth = Math.round(targetHeight * originalAspectRatio);
        }

        resizeWidth = Math.max(1, Math.min(resizeWidth, targetWidth));
        resizeHeight = Math.max(1, Math.min(resizeHeight, targetHeight));

        this.logger.debug(`区域 ${regionIndex}: 调整尺寸 ${metadata.width}x${metadata.height} -> ${resizeWidth}x${resizeHeight}`);

        // 计算居中的偏移量
        const offsetX = Math.floor((targetWidth - resizeWidth) / 2);
        const offsetY = Math.floor((targetHeight - resizeHeight) / 2);

        this.logger.debug(`区域 ${regionIndex}: 居中偏移 X=${offsetX}, Y=${offsetY}`);

        // 先调整大小并居中
        let resizedBuffer = await sharp(buffer)
            .resize(resizeWidth, resizeHeight, {
                fit: 'contain',
                background: { r: 255, g: 255, b: 255 }
            })
            .extend({
                top: offsetY,
                bottom: targetHeight - resizeHeight - offsetY,
                left: offsetX,
                right: targetWidth - resizeWidth - offsetX,
                background: { r: 255, g: 255, b: 255 }
            })
            .png()
            .toBuffer();

        // 在左右各添加10像素空白
        const finalWidth = targetWidth + 20; // 左右各加10像素
        const finalHeight = targetHeight;

        resizedBuffer = await sharp(resizedBuffer)
            .extend({
                top: 0,
                bottom: 0,
                left: 10,
                right: 10,
                background: { r: 255, g: 255, b: 255 }
            })
            .png()
            .toBuffer();

        this.logger.debug(`区域 ${regionIndex}: 最终尺寸 ${finalWidth}x${finalHeight} (左右各加10像素空白)`);

        return resizedBuffer;
    }

    async bufferToTensor(buffer, width, height) {
        // 获取实际图像尺寸（因为现在宽度增加了20像素）
        const metadata = await sharp(buffer).metadata();
        const actualWidth = metadata.width;
        const actualHeight = metadata.height;

        const imageData = await sharp(buffer)
            .ensureAlpha()
            .raw()
            .toBuffer({ resolveWithObject: true });

        // 使用实际尺寸创建张量
        const inputData = new Float32Array(3 * actualHeight * actualWidth);
        const data = imageData.data;

        for (let i = 0; i < data.length; i += 4) {
            const pixelIndex = Math.floor(i / 4);
            const y = Math.floor(pixelIndex / actualWidth);
            const x = pixelIndex % actualWidth;

            // 使用灰度值填充三个通道
            const grayValue = data[i] / 255.0;

            for (let c = 0; c < 3; c++) {
                const inputIndex = c * actualHeight * actualWidth + y * actualWidth + x;
                if (inputIndex < inputData.length) {
                    inputData[inputIndex] = grayValue;
                }
            }
        }

        return inputData;
    }

    postprocessRecognition(outputs) {
        this.logger.debug('开始识别后处理');

        try {
            const outputNames = this.recSession.outputNames;
            const recognitionOutput = outputs[outputNames[0]];

            if (!recognitionOutput) {
                this.logger.debug('识别输出为空');
                return { text: '', confidence: 0 };
            }

            const data = recognitionOutput.data;
            const [batch, seqLen, vocabSize] = recognitionOutput.dims;

            this.logger.debug(`序列长度: ${seqLen}, 词汇表大小: ${vocabSize}, 字符集大小: ${this.characterSet.length}`);

            if (this.characterSet.length === 0) {
                this.logger.error('字符集为空');
                return { text: '', confidence: 0 };
            }

            // 验证词汇表大小与字符集大小的匹配
            if (vocabSize !== this.characterSet.length + 1) {
                this.logger.warn(`词汇表大小(${vocabSize})与字符集大小(${this.characterSet.length})不匹配，可能影响识别效果`);
            }

            const { text, confidence } = this.ctcDecode(data, seqLen, vocabSize);
            this.logger.debug(`解码结果: "${text}", 置信度: ${confidence.toFixed(4)}`);

            return { text, confidence };

        } catch (error) {
            this.logger.error('识别后处理失败', error);
            return { text: '', confidence: 0 };
        }
    }

    ctcDecode(data, seqLen, vocabSize) {
        let text = '';
        let lastCharIndex = -1;
        let confidenceSum = 0;
        let charCount = 0;

        // 动态阈值调整
        const baseThreshold = 0.03;
        let confidenceThreshold = baseThreshold;

        // 分析序列置信度分布
        let maxSequenceProb = 0;
        let minSequenceProb = 1;
        let sumProb = 0;
        let probCount = 0;

        for (let t = 0; t < seqLen; t++) {
            for (let i = 0; i < vocabSize; i++) {
                const prob = data[t * vocabSize + i];
                if (prob > 0.01) { // 只统计有意义的概率
                    maxSequenceProb = Math.max(maxSequenceProb, prob);
                    minSequenceProb = Math.min(minSequenceProb, prob);
                    sumProb += prob;
                    probCount++;
                }
            }
        }

        const avgProb = probCount > 0 ? sumProb / probCount : 0;

        // 根据序列特性动态调整阈值
        if (avgProb < 0.3) {
            confidenceThreshold = baseThreshold * 0.5;
        } else if (avgProb > 0.7) {
            confidenceThreshold = baseThreshold * 1.5;
        }

        this.logger.debug(`序列统计: 平均概率=${avgProb.toFixed(4)}, 使用解码阈值: ${confidenceThreshold.toFixed(4)}`);

        // 改进的beam search算法
        const beamWidth = 5;
        let beams = [{ text: '', confidence: 1.0, lastChar: -1 }];

        for (let t = 0; t < seqLen; t++) {
            const newBeams = [];

            // 获取当前时间步的top-k字符
            const topK = [];
            for (let i = 0; i < vocabSize; i++) {
                const prob = data[t * vocabSize + i];
                if (prob > confidenceThreshold) {
                    topK.push({ index: i, prob });
                }
            }

            // 按概率排序
            topK.sort((a, b) => b.prob - a.prob);
            const candidates = topK.slice(0, beamWidth);

            // 为每个beam扩展候选字符
            for (const beam of beams) {
                for (const candidate of candidates) {
                    const charIndex = candidate.index;

                    if (charIndex === 0) {
                        // 空白字符
                        newBeams.push({
                            text: beam.text,
                            confidence: beam.confidence,
                            lastChar: -1
                        });
                    } else {
                        const actualCharIndex = charIndex - 1;
                        if (actualCharIndex < this.characterSet.length) {
                            const char = this.characterSet[actualCharIndex];
                            let newText = beam.text;

                            // 处理重复字符
                            if (charIndex !== beam.lastChar) {
                                newText += char;
                            }

                            newBeams.push({
                                text: newText,
                                confidence: beam.confidence * candidate.prob,
                                lastChar: charIndex
                            });
                        }
                    }
                }
            }

            // 选择top beamWidth个beam
            newBeams.sort((a, b) => b.confidence - a.confidence);
            beams = newBeams.slice(0, beamWidth);
        }

        // 选择最佳beam
        if (beams.length > 0) {
            const bestBeam = beams[0];
            text = bestBeam.text;

            // 计算平均置信度（几何平均）
            const textLength = text.length;
            if (textLength > 0) {
                confidenceSum = Math.pow(bestBeam.confidence, 1 / textLength);
                charCount = textLength;
            }
        }

        const avgConfidence = charCount > 0 ? confidenceSum : 0;

        return {
            text: text,
            confidence: avgConfidence
        };
    }

}

export default TextRecognizer;