from fastapi import FastAPI
from pydantic import BaseModel
import easyocr
import numpy as np
from PIL import Image
import cv2
import base64
import io
import os

app = FastAPI(title="OCR Base64 API")

model_storage_dir = '/code/.EasyOCR'
os.makedirs(model_storage_dir, exist_ok=True)

reader = easyocr.Reader(['en'], gpu=False, model_storage_directory=model_storage_dir)

class ImageRequest(BaseModel):
    image_base64: str

def preprocess_image(image):
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    scale_percent = 200
    width = int(gray.shape[1] * scale_percent / 100)
    height = int(gray.shape[0] * scale_percent / 100)
    resized = cv2.resize(gray, (width, height), interpolation=cv2.INTER_LINEAR)
    _, thresh = cv2.threshold(resized, 150, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    return thresh

@app.post("/extract-text/")
def extract_text(request: ImageRequest):
    try:
        b64_data = request.image_base64
        if "base64," in b64_data:
            b64_data = b64_data.split("base64,")[1]

        image_bytes = base64.b64decode(b64_data)
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        img_np = np.array(image)

        processed_img = preprocess_image(img_np)

        results = reader.readtext(processed_img)
        extracted_text = [text.replace(" ", "") for _, text, _ in results]
        
        text_result = " ".join(extracted_text) if extracted_text else ""

        return {
            "success": True,
            "text": text_result
        }

    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }

@app.get("/")
def read_root():
    return {"message": "API is running!"}