import json
from base64 import b64decode

import cv2
import numpy as np
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from paddleocr import PaddleOCR, PPStructure
from sx_utils.sxweb import *
from sx_utils.sximage import *

import os

# 初始化app
app = FastAPI()
origins = ["*"]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

use_gpu = os.getenv('USE_CUDA') == 'gpu'
print(f'use gpu: {use_gpu}')

# 普通表格
table_engine = PPStructure(layout=False,
                           table=True,
                           use_gpu=use_gpu,
                           show_log=True,
                           det_model_dir="models/det/det_table_v2",
                           rec_model_dir="./models/rec/rec_table_v1",
                           table_model_dir="models/table/SLAnet_v1")

# 长度较长表格
table_engine1 = PPStructure(layout=False,
                            table=True,
                            use_gpu=use_gpu,
                            show_log=True,
                            det_model_dir="models/det/det_table_v1",
                            rec_model_dir="./models/rec/rec_table_v1",
                            table_model_dir="./models/table/SLAnet_v1")

# 针对某些特殊情况的补充模型
table_engine2 = PPStructure(layout=False,
                            table=True,
                            use_gpu=use_gpu,
                            show_log=True,
                            det_model_dir="models/det/det_table_v3",
                            rec_model_dir="./models/rec/rec_table_v1",
                            table_model_dir="./models/table/SLAnet_v1")


class TableInfo(BaseModel):
    image: str
    det: str


@app.get("/ping")
def ping():
    return 'pong!'


@app.post("/ocr_system/table")
@web_try()
def table(image: TableInfo):
    img = base64_to_np(image.image)
    if image.det == 'no':
        res = table_engine(img)
    elif image.det == 'yes':
        res = table_engine1(img)
    elif image.det == 'spe':
        res = table_engine2(img)
    return res[0]['res']


if __name__ == '__main__':
    import uvicorn
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--host', default='0.0.0.0')
    parser.add_argument('--port', default=8080)
    opt = parser.parse_args()

    app_str = 'server:app'  # make the app string equal to whatever the name of this file is
    uvicorn.run(app_str, host=opt.host, port=int(opt.port), reload=True)