import io
import os
import tempfile
from fastapi import FastAPI, File, UploadFile, HTTPException, Query
from PIL import Image
import pytesseract
import pypdfium2 as pdfium

app = FastAPI(title="Scholarform OCR Service")

@app.get("/")
def root():
    return {"ok": True, "service": "ocr"}

@app.get("/health")
def health():
    return {"ok": True, "service": "ocr"}

def ocr_image_bytes(data: bytes, lang: str) -> str:
    image = Image.open(io.BytesIO(data)).convert("RGB")
    return pytesseract.image_to_string(image, lang=lang)

def ocr_pdf_file(path: str, lang: str, max_pages: int) -> str:
    pdf = pdfium.PdfDocument(path)
    pages = min(len(pdf), max_pages)
    texts = []
    for i in range(pages):
        page = pdf[i]
        pil_img = page.render(scale=2).to_pil()
        texts.append(pytesseract.image_to_string(pil_img, lang=lang))
    return "\n".join(texts)

@app.post("/ocr")
async def ocr(
    file: UploadFile = File(...),
    lang: str = Query("eng"),
    max_pages: int = Query(10, ge=1, le=200),
):
    name = (file.filename or "").lower()
    content = await file.read()

    if not content:
        raise HTTPException(status_code=400, detail="Empty file")

    try:
        if name.endswith(".pdf"):
            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
                tmp.write(content)
                path = tmp.name
            try:
                text = ocr_pdf_file(path, lang=lang, max_pages=max_pages)
            finally:
                try:
                    os.remove(path)
                except Exception:
                    pass
            return {"ok": True, "type": "pdf", "text": text}

        if name.endswith((".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff")):
            text = ocr_image_bytes(content, lang=lang)
            return {"ok": True, "type": "image", "text": text}

        raise HTTPException(status_code=400, detail="Supported: PDF or image")
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"OCR failed: {e}")