plugins/ocr/api.py

"""OCR 插件 HTTP 接口。

所有识别请求走 FastAPI BackgroundTasks 真正异步执行：上传/解码完成后立即返回
run_id，OCR 在后台任务中完成并通过 run/log 服务写事件，前端轮询日志面板即可。
"""
import base64
import logging
import tempfile
from pathlib import Path

from fastapi import APIRouter, BackgroundTasks, HTTPException, UploadFile
from pydantic import BaseModel

from .core import extract_text, get_model_info_with_providers
from app.plugins.run_log import get_run_log_service, RunStatus

logger = logging.getLogger(__name__)
router = APIRouter()
plugin = None


def set_plugin_instance(plugin_instance):
    global plugin
    plugin = plugin_instance


@router.get("/status")
async def get_status():
    """三态状态：插件启用 / 引擎就绪 / 模型 ID。"""
    if plugin is None:
        return {"name": "ocr", "enabled": False, "engine_ready": False}
    return plugin.get_status()


@router.get("/model-info")
async def model_info():
    return get_model_info_with_providers()


class OCRBase64Request(BaseModel):
    image_base64: str
    suffix: str = ".jpg"


@router.post("/upload")
async def ocr_upload(file: UploadFile, background: BackgroundTasks):
    """上传图片并异步识别，立即返回 run_id。"""
    run = _create_run(f"图片已上传: {file.filename}")
    content = await file.read()
    background.add_task(
        _run_ocr_task, content, _suffix_from_filename(file.filename), run.run_id,
    )
    return {"run_id": run.run_id, "status": "running"}


@router.post("/image-base64")
async def ocr_image_base64(request: OCRBase64Request, background: BackgroundTasks):
    """Base64 图片异步识别。"""
    try:
        content = base64.b64decode(request.image_base64)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Base64 解码失败: {e}")
    run = _create_run("Base64 图片已接收")
    background.add_task(
        _run_ocr_task, content, _normalize_suffix(request.suffix), run.run_id,
    )
    return {"run_id": run.run_id, "status": "running"}


# ---- 后台任务 ----

def _create_run(upload_msg: str):
    """创建 run 并写上传事件；前置校验 plugin 是否启用。"""
    if plugin is None or not plugin.enabled:
        raise HTTPException(status_code=400, detail="插件未启用")
    run_service = get_run_log_service()
    run = run_service.create_run("ocr")
    run_service.add_event(run_id=run.run_id, stage="upload", message=upload_msg)
    return run


def _run_ocr_task(content: bytes, suffix: str, run_id: str) -> None:
    """后台识别任务：落临时文件 → 调引擎 → 写 run 结果 → 清理临时文件。

    整段 try/finally 保证临时文件最终被清理；run_service.finish_run 一定被调用。
    """
    run_service = get_run_log_service()
    tmp_path: str | None = None
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            tmp.write(content)
            tmp_path = tmp.name
        result = extract_text(tmp_path, run_id=run_id)
        if result["success"]:
            run_service.finish_run(
                run_id=run_id,
                status=RunStatus.SUCCEEDED,
                result={
                    "text": result["text"],
                    "lines": result["lines"],
                    "boxes": result["boxes"],
                    "model_info": result["model_info"],
                    "timings": result["timings"],
                },
            )
        else:
            run_service.finish_run(
                run_id=run_id, status=RunStatus.FAILED, error=result["error"],
            )
    except Exception as e:
        logger.exception("后台 OCR 异常")
        run_service.add_event(run_id, "ocr", f"OCR 异常: {e}", level="error")
        run_service.finish_run(run_id=run_id, status=RunStatus.FAILED, error=str(e))
    finally:
        if tmp_path:
            Path(tmp_path).unlink(missing_ok=True)


# ---- 工具 ----

def _suffix_from_filename(filename: str) -> str:
    return "." + filename.rsplit(".", 1)[-1] if "." in filename else ".jpg"


def _normalize_suffix(suffix: str) -> str:
    return suffix if suffix.startswith(".") else f".{suffix}"