import uvicorn
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from typing import Optional
import numpy as np
from PIL import Image
from paddleocr import PaddleOCR
from doctr.io import DocumentFile
from doctr.models import ocr_predictor
import io
import logging

class OCRAPIApp:
    def __init__(self):
        self.app = FastAPI(
            docs_url="/",
            title="OCR API",
            version="1.0",
        )
        self.setup_routes()
        self.paddle_ocr = PaddleOCR(lang='en', use_angle_cls=True)
        self.doctr_model = ocr_predictor(pretrained=True)

    def ocr_with_paddle(self, img):
        try:
            logging.info("Processing image with PaddleOCR...")
            result = self.paddle_ocr.ocr(img)
            text_output = ' '.join([line[1][0] for line in result[0]])
            return text_output
        except Exception as e:
            logging.error(f"Error with PaddleOCR: {e}")
            raise HTTPException(status_code=500, detail="Error processing image")

    def ocr_with_doctr(self, file):
        try:
            logging.info("Processing PDF with Doctr...")
            doc = DocumentFile.from_pdf(file)
            result = self.doctr_model(doc)
            text_output = ''
            for page in result.pages:
                for block in page.blocks:
                    for line in block.lines:
                        text_output += ' '.join([word.value for word in line.words]) + "\n"
            return text_output
        except Exception as e:
            logging.error(f"Error with Doctr: {e}")
            raise HTTPException(status_code=500, detail="Error processing PDF")

    async def ocr_endpoint(self, file: UploadFile = File(...)):
        try:
            file_bytes = await file.read()
            if file.filename.endswith(".pdf"):
                text_output = self.ocr_with_doctr(io.BytesIO(file_bytes))
            else:
                img = np.array(Image.open(io.BytesIO(file_bytes)))
                text_output = self.ocr_with_paddle(img)
            return {"ocr_text": text_output}
        except Exception as e:
            logging.error(f"Error processing file: {e}")
            raise HTTPException(status_code=500, detail="Error processing file")

    def setup_routes(self):
        self.app.post("/ocr")(self.ocr_endpoint)

# Initialize the app
app = OCRAPIApp().app

# Add CORS middleware for cross-origin requests
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
if __name__ == "__main__":
   python -m app