from fastapi import FastAPI
from pydantic import BaseModel
import base64
import subprocess
import re
import uuid
import os
import logging
from typing import List

app = FastAPI()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

ansi_escape = re.compile(r'\x1b[^m]*m')


class PictureRequest(BaseModel):
    pic_id: int
    pic: str


class PictureResponse(BaseModel):
    pic_id: int
    pic_txt: str
    code: int


@app.post("/ocr", response_model=List[PictureResponse])
async def ocr_endpoint(pictures: List[PictureRequest]):
    results = []

    for picture in pictures:
        image_data = base64.b64decode(picture.pic)
        temp_image_file = f"{uuid.uuid4()}.jpg"  # 使用uuid生成唯一的临时文件名

        try:
            with open(temp_image_file, "wb") as f:
                f.write(image_data)

            command = [
                'python', 'tools/infer/predict_system_1.py',
                '--use_gpu=False',
                '--cls_model_dir=./models/cls',
                '--rec_model_dir=./models/rec',
                '--det_model_dir=./models/det',
                f'--image_dir={temp_image_file}'
            ]

            logger.info(f"Processing image with pic_id: {picture.pic_id}")
            result = subprocess.run(command, capture_output=True, text=True, encoding='utf-8', errors='ignore')

            if result.returncode == 0:
                ocr_str = ansi_escape.sub('', result.stdout).strip()
                code = 200
                logger.info(f"OCR successful for pic_id: {picture.pic_id}")
            else:
                ocr_str = ""
                code = 500
                logger.error(f"OCR failed for pic_id: {picture.pic_id}, return code: {result.returncode}")

        except Exception as e:
            ocr_str = str(e)
            code = 500
            logger.exception(f"Exception occurred while processing pic_id: {picture.pic_id}")

        finally:
            if os.path.exists(temp_image_file):
                os.remove(temp_image_file)  # 确保临时文件在处理完成后被删除

        results.append({"pic_id": picture.pic_id, "pic_txt": ocr_str, "code": code})

    return results


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)