app.py

import os
import hashlib
import tempfile
import threading
import traceback
from collections import OrderedDict

from fastapi import FastAPI, UploadFile, File, HTTPException, Request
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
import wcocr

app = FastAPI(title="WeChat OCR API", version="1.0.0")

# 全局变量
ocr_cache = None
ocr_lock = None

API_TOKEN = os.getenv("TOKEN", "my-secret-token")

# 不需要验证的路径
EXCLUDED_PATHS = ["/health"]


class TokenAuthMiddleware(BaseHTTPMiddleware):
    """简单的Token验证中间件"""

    async def dispatch(self, request: Request, call_next):
        # 跳过不需要验证的路径
        if request.url.path in EXCLUDED_PATHS:
            return await call_next(request)

        # 从请求头获取 Authorization
        auth_header = request.headers.get("Authorization")

        # 兼容 Bearer 前缀
        if auth_header and auth_header.startswith("Bearer "):
            auth_header = auth_header[7:]  # 移除 "Bearer " 前缀

        # 验证token是否匹配
        if auth_header != API_TOKEN:
            return JSONResponse(
                status_code=401,
                content={"detail": "未授权：无效的认证令牌"}
            )

        # token验证通过，继续处理请求
        response = await call_next(request)
        return response


# 添加中间件
app.add_middleware(TokenAuthMiddleware)


class LRUCache:
    """简单的 LRU 缓存，防止内存溢出"""

    def __init__(self, capacity: int = 100):
        self.cache = OrderedDict()
        self.capacity = capacity

    def get(self, key):
        if key not in self.cache:
            return None
        self.cache.move_to_end(key)
        return self.cache[key]

    def put(self, key, value):
        if key in self.cache:
            self.cache.move_to_end(key)
        self.cache[key] = value
        if len(self.cache) > self.capacity:
            self.cache.popitem(last=False)


@app.on_event("startup")
async def startup_event():
    global ocr_cache, ocr_lock

    print("====== 开始初始化 WeChat OCR ======")
    try:
        # 初始化 WeChat OCR
        WECHAT_PATH = os.environ.get("WECHAT_PATH", "/code/wechat")
        WXOCR_PATH = os.environ.get("WXOCR_PATH", "/code/wechat/wxocr")
        wcocr.init(WXOCR_PATH, WECHAT_PATH)

        # 全局缓存对象 (缓存 100 张最近的图片结果)
        ocr_cache = LRUCache(capacity=100)

        # OCR 引擎并发锁
        ocr_lock = threading.Lock()

        print("====== 🎉 WeChat OCR 初始化成功，API 已就绪！ ======")
    except Exception as e:
        print("❌ 初始化失败, 错误如下:\n", str(e))


@app.get("/health")
def health_check():
    return {"status": "healthy"}


def process_ocr(file_bytes: bytes) -> dict:
    """处理图片并执行 OCR，包含缓存机制"""
    # 计算图片 Hash 用于缓存 Key
    img_hash = hashlib.md5(file_bytes).hexdigest()
    cached_result = ocr_cache.get(img_hash) if ocr_cache else None
    
    if cached_result is not None:
        return cached_result

    if not file_bytes or len(file_bytes) < 8:
        return {"errcode": 400, "errmsg": "Empty or invalid file content"}

    # 使用轻量级的 Magic Bytes（文件头魔数）检查代替 cv2，判断是否是常见图片格式
    header = file_bytes[:12]
    is_image = (
        header.startswith(b'\x89PNG\r\n\x1a\n') or  # PNG
        header.startswith(b'\xff\xd8\xff') or      # JPEG
        header.startswith(b'BM') or                # BMP
        header.startswith(b'GIF8') or              # GIF
        (header.startswith(b'RIFF') and header[8:12] == b'WEBP') # WEBP
    )
    
    if not is_image:
        return {"errcode": 400, "errmsg": "Unsupported or invalid image format"}

    # 写入临时文件供 wcocr 调用
    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_file:
        tmp_file.write(file_bytes)
        tmp_file.flush()
        tmp_path = tmp_file.name

    try:
        # 加锁执行 OCR，防止底层组件崩溃
        if ocr_lock:
            with ocr_lock:
                result = wcocr.ocr(tmp_path)
        else:
            result = wcocr.ocr(tmp_path)
    finally:
        # 清理临时文件
        if os.path.exists(tmp_path):
            os.remove(tmp_path)

    # 存入缓存
    if ocr_cache:
        ocr_cache.put(img_hash, result)
    return result


def format_structured_text(ocr_response: list) -> str:
    """基于坐标的文本行合并算法，还原原始文本的换行结构"""
    if not ocr_response:
        return ""

    # 按 Y 轴 (top) 排序
    boxes = sorted(ocr_response, key=lambda x: x['top'])
    lines = []
    current_line = []

    for box in boxes:
        if not current_line:
            current_line.append(box)
            continue

        # 计算当前行的中心基线
        line_top_avg = sum(b['top'] for b in current_line) / len(current_line)
        line_bottom_avg = sum(b['bottom']
                              for b in current_line) / len(current_line)
        line_center = (line_top_avg + line_bottom_avg) / 2

        # 如果当前 box 的 top 小于上一行的中心点，认为在同一行
        if box['top'] < line_center:
            current_line.append(box)
        else:
            lines.append(current_line)
            current_line = [box]

    if current_line:
        lines.append(current_line)

    # 将同一行的文本按 X 轴 (left) 排序并拼接
    structured_text = []
    for line in lines:
        line.sort(key=lambda x: x['left'])
        line_text = " ".join([b['text'] for b in line])
        structured_text.append(line_text)

    return "\n".join(structured_text)


# ==========================================
# 接口 1：OCR 识别后的原始结构
# ==========================================

@app.post("/api/ocr/raw")
async def api_ocr_raw(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        raw_result = process_ocr(contents)

        if raw_result.get('errcode', 0) != 0:
            return JSONResponse(status_code=400, content=raw_result)

        return raw_result
    except Exception as e:
        print(f"❌ [api_ocr_raw] 发生异常: {str(e)}")
        traceback.print_exc()
        raise HTTPException(
            status_code=500, detail=f"OCR raw extraction error: {str(e)}")


# ==========================================
# 接口 2：识别后的纯文本 (无分隔符直接拼接)
# ==========================================

@app.post("/api/ocr/text")
async def api_ocr_text(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        raw_result = process_ocr(contents)

        if raw_result.get('errcode', 0) != 0:
            return JSONResponse(status_code=400, content=raw_result)

        ocr_response = raw_result.get('ocr_response', [])
        pure_text = "".join([item['text'] for item in ocr_response])

        return {
            "errcode": 0,
            "text": pure_text
        }
    except Exception as e:
        print(f"❌ [api_ocr_text] 发生异常: {str(e)}")
        traceback.print_exc()
        raise HTTPException(
            status_code=500, detail=f"OCR text extraction error: {str(e)}")


# ==========================================
# 接口 3：识别后的纯文本，保留原始结构换行
# ==========================================

@app.post("/api/ocr/structured")
async def api_ocr_structured(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        raw_result = process_ocr(contents)

        if raw_result.get('errcode', 0) != 0:
            return JSONResponse(status_code=400, content=raw_result)

        ocr_response = raw_result.get('ocr_response', [])
        structured_text = format_structured_text(ocr_response)

        return {
            "errcode": 0,
            "text": structured_text
        }
    except Exception as e:
        print(f"❌ [api_ocr_structured] 发生异常: {str(e)}")
        traceback.print_exc()
        raise HTTPException(
            status_code=500, detail=f"OCR structured extraction error: {str(e)}")