import os
import shutil
import uuid
import json
from flask import Flask, request, jsonify
import paddle
from paddleocr import PaddleOCR
import gc
import time
# ---------------- 配置 ----------------
os.environ["METAX_GPU_VISIBLE_DEVICES"] = "2"  # 适配沐曦C500显卡，指定可用设备
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
UPLOAD_DIR = os.path.join(BASE_DIR, "upload_images")
OUT_DIR = os.path.join(BASE_DIR, "ocr_output")
GPU_ID = "metax_gpu:2"  # 适配沐曦C500显卡，使用metax_gpu设备

os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(OUT_DIR, exist_ok=True)

# ---------------- 初始化 OCR 模型 ----------------
# 优化显存使用，对于64GB显存可以使用更大比例
os.environ['FLAGS_fraction_of_gpu_memory_to_use'] = '0.8'  # 增加显存使用比例
os.environ['FLAGS_cudnn_deterministic'] = '0'  # 关闭确定性算法，提高性能
os.environ['FLAGS_use_stream_safe_cuda_allocator'] = '0'  # 关闭流安全分配器，提高性能

paddle.set_device(GPU_ID)

print(f"[init-{GPU_ID}] 正在加载模型...")
# 优化模型加载，保持必要功能的同时提高启动速度
ocr = PaddleOCR(
    device=GPU_ID,
    use_doc_orientation_classify=True,  # 保持文档方向分类功能
    use_doc_unwarping=True,  # 保持文档校正功能
    use_textline_orientation=True,  # 保持文本行方向分类功能
    enable_mkldnn=False,  # 禁用MKL-DNN，专注于GPU计算
    cpu_threads=1  # 减少CPU线程使用，专注于GPU
)
# ocr1 = PaddleOCR(
#     device="metax_gpu:1",  # 如需使用多卡，可配置metax_gpu:1
#     use_doc_orientation_classify=True,  # 保持文档方向分类功能
#     use_doc_unwarping=True,  # 保持文档校正功能
#     use_textline_orientation=True,  # 保持文本行方向分类功能
#     enable_mkldnn=False,  # 禁用MKL-DNN，专注于GPU计算
#     cpu_threads=1  # 减少CPU线程使用，专注于GPU
# )
print(f"[init-{GPU_ID}] 模型加载完成 ✅")

# ---------------- Flask 应用 ----------------
app = Flask(__name__)


@app.route("/ocr", methods=["POST"])
def ocr_endpoint():


    eab = None
    start_time = time.perf_counter()
    if "files" not in request.files:
        return jsonify({"error": "请上传图片"}), 400

    files = request.files.getlist("files")
    if not files:
        return jsonify({"error": "未检测到文件"}), 400

    session_id = str(uuid.uuid4())[:8]
    session_dir = os.path.join(UPLOAD_DIR, session_id)
    os.makedirs(session_dir, exist_ok=True)

    results_all = []

    for file_idx, file in enumerate(files):
        filename = file.filename
        img_path = os.path.join(session_dir, filename)
        file.save(img_path)

        out_path = os.path.join(OUT_DIR, f"{filename}.json")

        try:
            # 直接进行OCR预测，不保存中间文件
            result = ocr.predict(input=img_path)
            
            # 直接从OCR结果中提取文本，避免JSON文件I/O操作
            eab = ""
            for res in result:
                try:
                    # 直接从结果对象中提取文本
                    if hasattr(res, 'boxes'):
                        for box, text, score in res.boxes:
                            eab += text
                    elif hasattr(res, 'rec_texts'):
                        eab = ''.join(res.rec_texts)
                    else:
                        # 尝试获取文本的其他方式
                        res.print()
                        # 保存到JSON作为备份
                        res.save_to_json(out_path)
                        with open(out_path, "r", encoding="utf-8") as f:
                            data = json.load(f)
                        eab = ''.join(data.get('rec_texts', []))
                except Exception as save_err:
                    print(f"保存结果失败: {save_err}")

        except Exception as e:
            results_all.append({
                "error": str(e)
            })
        finally:
            # 仅在必要时释放内存，减少频繁内存操作
            if file_idx % 5 == 0:  # 每处理5张图片释放一次内存
                paddle.device.empty_cache()
                gc.collect()
    end_time = time.perf_counter()  # ⏱️ 结束计时
    elapsed_time = round(end_time - start_time, 2)  # 用时保留 2 位小数（秒）
    return jsonify({
        "time": f"{elapsed_time}s",
        "results":eab
    })


@app.route("/clear", methods=["POST"])
def clear_output():
    shutil.rmtree(UPLOAD_DIR, ignore_errors=True)
    shutil.rmtree(OUT_DIR, ignore_errors=True)
    os.makedirs(UPLOAD_DIR, exist_ok=True)
    os.makedirs(OUT_DIR, exist_ok=True)
    return jsonify({"status": "cleared"})


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5001)