from fastapi import FastAPI, UploadFile, File, HTTPException
from paddleocr import PaddleOCR
from PIL import Image
import io, numpy as np
import fitz  # pymupdf

app = FastAPI()
ocr = PaddleOCR(use_angle_cls=True, lang='en', use_gpu=False)

@app.get("/")
def health():
    return {"status": "ok"}

@app.post("/ocr")
async def run_ocr(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        all_texts = []

        # Handle PDF
        if file.filename.lower().endswith(".pdf"):
            pdf = fitz.open(stream=contents, filetype="pdf")
            for page in pdf:
                pix = page.get_pixmap(dpi=200)
                img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                img_array = np.array(img)
                result = ocr.ocr(img_array, cls=True)
                if result and result[0]:
                    for line in result[0]:
                        all_texts.append({
                            "text": line[1][0],
                            "confidence": round(line[1][1], 4)
                        })

        # Handle image (JPG, PNG, etc)
        else:
            image = Image.open(io.BytesIO(contents)).convert("RGB")
            img_array = np.array(image)
            result = ocr.ocr(img_array, cls=True)
            if result and result[0]:
                for line in result[0]:
                    all_texts.append({
                        "text": line[1][0],
                        "confidence": round(line[1][1], 4)
                    })

        return {"filename": file.filename, "results": all_texts}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))