from fastapi import FastAPI, UploadFile, File
from paddleocr import PPStructure
from PIL import Image, ImageEnhance, ImageFilter
import numpy as np, cv2, io

app = FastAPI()

# Loaded globally, but memory footprint is constrained by HF Free Tier limits
table_engine = PPStructure(show_log=False, image_orientation=True, lang='en')

@app.get("/")
def home(): return {"status": "Production Engine Ready"}

def _deskew(image_np):
    gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)
    lines = cv2.HoughLinesP(edges, 1, np.pi/180, 100, minLineLength=100, maxLineGap=10)
    if lines is not None:
        angles = [np.degrees(np.arctan2(l[0][3]-l[0][1], l[0][2]-l[0][0])) for l in lines]
        valid_angles = [a for a in angles if -15 < a < 15]
        if valid_angles and abs(np.median(valid_angles)) > 0.5:
            h, w = image_np.shape[:2]
            M = cv2.getRotationMatrix2D((w//2, h//2), np.median(valid_angles), 1.0)
            return cv2.warpAffine(image_np, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_CONSTANT, borderValue=(255, 255, 255))
    return image_np

def _prep(img):
    img_np = np.array(img)
    # Optimization: Only run heavy deskew and contrast if it's a messy scan
    if cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY).std() < 55:
        img_np = _deskew(img_np)
        pil_img = ImageEnhance.Contrast(Image.fromarray(img_np).convert("L")).enhance(1.8).filter(ImageFilter.SHARPEN).convert("RGB")
        return np.array(pil_img)
    return img_np

def _sort(res):
    # FIXED: Properly handles PPStructure's nested tuple format
    b = []
    for item in res:
        if isinstance(item, list) and len(item) == 2 and isinstance(item[1], (tuple, list)):
            box, (text, conf) = item[0], item[1]
            if float(conf) >= 0.45:
                b.append({"text": text, "y": min(p[1] for p in box), "x": min(p[0] for p in box)})
    return "\n".join([x["text"] for x in sorted(b, key=lambda k: (round(k["y"]/15), k["x"]))])

@app.post("/ocr")
async def get_ocr(file: UploadFile = File(...)):
    try:
        res = table_engine(_prep(Image.open(io.BytesIO(await file.read())).convert("RGB")))
        text_blocks = []
        for r in res:
            if r["type"] == "table":
                text_blocks.append(f"[TABLE]\n{r.get('res', {}).get('html', '')}")
            else:
                sorted_txt = _sort(r.get("res", []))
                if sorted_txt: text_blocks.append(sorted_txt)
        return {"text": "\n\n".join(text_blocks)}
    except Exception as e:
        return {"text": "", "error": str(e)}