#WECHAT_OCR_DIR = "C:\\Users\\docker\\AppData\\Roaming\\Tencent\\WeChat\\XPlugin\\Plugins\\WeChatOCR\\7079\\extracted\\WeChatOCR.exe"
#WECHAT_DIR = "c:\\Users\\Docker\\WeChat\\[3.9.12.17]"
import os
import json
import time
import threading
import logging
from flask import Flask, request, jsonify
from werkzeug.utils import secure_filename
import requests
from io import BytesIO
from wechat_ocr.ocr_manager import OcrManager, OCR_MAX_TASK_ID

# 配置日志
logging.basicConfig(level=logging.DEBUG, 
                    format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 配置

WECHAT_OCR_DIR = "C:\\Users\\docker\\AppData\\Roaming\\Tencent\\WeChat\\XPlugin\\Plugins\\WeChatOCR\\7079\\extracted\\WeChatOCR.exe"
WECHAT_DIR = "c:\\Users\\Docker\\WeChat\\[3.9.12.17]"
UPLOAD_FOLDER = os.path.abspath('uploads')
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp'}

# 确保上传文件夹存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

class OCRService:
    _instance = None
    _lock = threading.Lock()

    def __new__(cls):
        if not cls._instance:
            with cls._lock:
                if not cls._instance:
                    cls._instance = super(OCRService, cls).__new__(cls)
                    cls._instance._initialize()
        return cls._instance

    def _initialize(self):
        logger.info("初始化OCR服务")
        self.ocr_manager = OcrManager(WECHAT_DIR)
        self.ocr_manager.SetExePath(WECHAT_OCR_DIR)
        self.ocr_manager.SetUsrLibDir(WECHAT_DIR)
        self.results = {}
        self.lock = threading.Lock()

    def ocr_result_callback(self, img_path: str, results: dict):
        logger.info(f"OCR识别回调 - 图片: {img_path}")
        with self.lock:
            self.results[img_path] = results
            logger.debug(f"识别结果: {results}")

    def process_ocr(self, image_path):
        logger.info(f"开始处理图片: {image_path}")
        
        # 重置结果
        with self.lock:
            self.results = {}

        # 设置回调函数
        self.ocr_manager.SetOcrResultCallback(self.ocr_result_callback)

        # 启动OCR服务
        try:
            logger.info("启动WeChat OCR服务")
            self.ocr_manager.StartWeChatOCR()
            logger.info("OCR服务启动成功")

            # 执行OCR任务
            logger.info(f"开始识别图片: {image_path}")
            self.ocr_manager.DoOCRTask(image_path)
            logger.info("OCR任务提交成功")

            # 等待任务完成
            start_time = time.time()
            while self.ocr_manager.m_task_id.qsize() != OCR_MAX_TASK_ID:
                if time.time() - start_time > 30:  # 30秒超时
                    logger.error("OCR处理超时")
                    self.ocr_manager.KillWeChatOCR()
                    return {"error": "OCR processing timed out"}
                time.sleep(0.1)

            # 获取结果
            with self.lock:
                result = self.results.get(image_path, {})
                logger.info(f"获取OCR识别结果: {result}")

            return result

        except Exception as e:
            logger.error(f"OCR处理发生错误: {str(e)}")
            return {"error": str(e)}
        finally:
            # 确保关闭OCR服务
            logger.info("关闭WeChat OCR服务")
            self.ocr_manager.KillWeChatOCR()

# Flask 应用
app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB 文件大小限制

# OCR服务单例
ocr_service = OCRService()

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@app.route('/ocr', methods=['POST'])
def ocr_endpoint():
    logger.info("收到OCR请求")
    filepath = None
    try:
        # 处理文件上传
        if 'file' in request.files:
            file = request.files['file']
            
            if not file or file.filename == '':
                logger.error("没有选择文件")
                return jsonify({"error": "没有选择文件"}), 400
            
            if not allowed_file(file.filename):
                logger.error("不支持的文件类型")
                return jsonify({"error": "不支持的文件类型"}), 400
            
            filename = secure_filename(file.filename)
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(filepath)
            logger.info(f"文件保存成功: {filepath}")
        
        # 处理URL图片
        elif request.is_json and 'image_url' in request.json:
            url = request.json['image_url']
            logger.info(f"处理图片URL: {url}")
            response = requests.get(url)
            image = BytesIO(response.content)
            filename = secure_filename(os.path.basename(url))
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            with open(filepath, 'wb') as f:
                f.write(image.getbuffer())
        
        else:
            logger.error("未提供图片")
            return jsonify({"error": "需要提供图片URL或文件"}), 400

        # 执行OCR识别
        result = ocr_service.process_ocr(filepath)
        
        # 删除临时文件
        if filepath and os.path.exists(filepath):
            os.remove(filepath)
            logger.info(f"删除临时文件: {filepath}")
        
        return jsonify({
            "result": result,
            "status": "success"
        })
    
    except Exception as e:
        # 确保即使出错也删除临时文件
        if filepath and os.path.exists(filepath):
            os.remove(filepath)
        
        logger.error(f"OCR处理失败: {str(e)}")
        return jsonify({
            "error": f"OCR处理失败: {str(e)}",
            "status": "error"
        }), 500
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)

# 使用说明:
# 1. 确保已安装所有依赖: flask, requests, wechat_ocr
# 2. 根据实际情况修改 WECHAT_OCR_DIR 和 WECHAT_DIR 路径
# 3. 确保 WeChat OCR 插件已正确安装