from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pix2text import Pix2Text
import io
from PIL import Image
import uvicorn
import os

# 高速トークナイザーの並列処理によるデッドロックを回避
os.environ["TOKENIZERS_PARALLELISM"] = "false"

app = FastAPI()

# CORSを許可
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Pix2Textを初期化
p2t = Pix2Text()

@app.get("/")
def read_root():
    return {"status": "ok", "message": "Pix2Text API is running"}

@app.get("/status")
def get_status():
    return {"status": "ready"}

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="File must be an image")
    
    try:
        contents = await file.read()
        image = Image.open(io.BytesIO(contents)).convert("RGB")
        
        # Pix2Textで認識 (formula認識のみを指定)
        res = p2t.recognize_formula(image)
        
        return {"latex": res}
    except Exception as e:
        return {"error": str(e)}

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