from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from PIL import Image
import pytesseract
from pix2text import Pix2Text
from paddleocr import PaddleOCR
import io

app = FastAPI()

# ✅ Allow all origins (development only)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize OCR engines
ocr = PaddleOCR(
    use_doc_orientation_classify=False,
    use_doc_unwarping=False,
    use_textline_orientation=False,
)

p2t = Pix2Text()


@app.post("/ocr")
async def extract_text(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        image = Image.open(io.BytesIO(contents)).convert("RGB")

        # Step 1: Extract general text (paddleocr better than tesseract for mixed fonts)
        result = ocr.predict(contents, cls=True)
        text_blocks = [" ".join([line[1][0] for line in res]) for res in result]

        # Step 2: Try detecting math expressions
        math_output = ""
        try:
            math_output = p2t(image)
        except Exception as e:
            math_output = "No math detected"

        return JSONResponse(
            content={
                "general_text": " ".join(text_blocks),
                "math_equations": math_output,
            }
        )

    except Exception as e:
        return JSONResponse(content={"error": str(e)}, status_code=500)