import os
import io
import easyocr
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
from PIL import Image

# =========================
# EasyOCR config
# =========================
MODEL_DIR = "/app/.EasyOCR"
USER_NET_DIR = os.path.join(MODEL_DIR, "user_network")

os.makedirs(MODEL_DIR, exist_ok=True)
os.makedirs(USER_NET_DIR, exist_ok=True)

# ✅ preload reader with cached models
reader = easyocr.Reader(
    ['en', 'hi'],  # langs (can reduce to ['en'] if you want smaller image)
    model_storage_directory=MODEL_DIR,
    user_network_directory=USER_NET_DIR,
    download_enabled=False   # 🚫 no runtime downloads
)

# =========================
# FastAPI app
# =========================
app = FastAPI()

@app.get("/")
async def root():
    return {"message": "OCR API is running on Hugging Face 🚀"}

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

        # run OCR
        results = reader.readtext(image)

        # format results
        text_results = []
        for bbox, text, prob in results:
            text_results.append({
                "bbox": bbox,
                "text": text,
                "confidence": float(prob)
            })

        return JSONResponse(content={"results": text_results})

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