"""
ocr_engine.py — Cloud version.
Uses trocr-base-handwritten (smaller, half RAM of large version)
and Donut medical model. Both run locally on HF Spaces hardware.
"""

from functools import lru_cache
from dataclasses import dataclass
from typing import List

import numpy as np
import cv2
import torch
from PIL import Image
from transformers import TrOCRProcessor, VisionEncoderDecoderModel

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# base instead of large — 300MB vs 1.3GB, fits HF free RAM
TROCR_MODEL_ID = "microsoft/trocr-base-handwritten"
DONUT_MODEL_ID = "chinmays18/medical-prescription-ocr"


@dataclass
class OCRResult:
    engine: str
    raw_text: str
    lines: List[str]


@lru_cache(maxsize=1)
def _load_trocr():
    processor = TrOCRProcessor.from_pretrained(TROCR_MODEL_ID)
    model = VisionEncoderDecoderModel.from_pretrained(TROCR_MODEL_ID)
    model.to(DEVICE)
    model.eval()
    return processor, model


@lru_cache(maxsize=1)
def _load_donut():
    from transformers import DonutProcessor, VisionEncoderDecoderModel as DonutModel
    processor = DonutProcessor.from_pretrained(DONUT_MODEL_ID)
    model = DonutModel.from_pretrained(DONUT_MODEL_ID)
    model.to(DEVICE)
    model.eval()
    return processor, model


def _get_line_boxes(image: Image.Image) -> List[tuple]:
    arr = np.array(image.convert("L"))
    _, thresh = cv2.threshold(arr, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (40, 3))
    dilated = cv2.dilate(thresh, kernel, iterations=2)
    contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    boxes = [cv2.boundingRect(c) for c in contours]
    boxes = [b for b in boxes if b[2] > 60 and b[3] > 10]
    boxes.sort(key=lambda b: b[1])
    return boxes


def _is_meaningful(text: str) -> bool:
    if len(text) < 3: return False
    alpha_ratio = sum(c.isalpha() for c in text) / max(len(text), 1)
    if alpha_ratio < 0.3: return False
    if all(c in "0123456789 ." for c in text): return False
    return True


def run_trocr(image: Image.Image) -> OCRResult:
    processor, model = _load_trocr()
    boxes = _get_line_boxes(image)
    np_img = np.array(image.convert("RGB"))
    lines, seen = [], set()

    for (x, y, w, h) in boxes:
        pad = 5
        crop = np_img[max(0,y-pad):min(np_img.shape[0],y+h+pad),
                      max(0,x-pad):min(np_img.shape[1],x+w+pad)]
        if crop.size == 0: continue
        pixel_values = processor(
            images=Image.fromarray(crop).convert("RGB"), return_tensors="pt"
        ).pixel_values.to(DEVICE)
        with torch.no_grad():
            generated_ids = model.generate(pixel_values, max_length=128, num_beams=3)
        text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
        if text and text not in seen and _is_meaningful(text):
            seen.add(text)
            lines.append(text)

    return OCRResult(engine="trocr-base", raw_text="\n".join(lines), lines=lines)


def run_donut(image: Image.Image) -> OCRResult:
    processor, model = _load_donut()
    pixel_values = processor(
        images=image.convert("RGB"), return_tensors="pt"
    ).pixel_values.to(DEVICE)
    task_prompt = "<s_ocr>"
    decoder_input_ids = processor.tokenizer(
        task_prompt, add_special_tokens=False, return_tensors="pt"
    ).input_ids.to(DEVICE)
    with torch.no_grad():
        generated_ids = model.generate(
            pixel_values, decoder_input_ids=decoder_input_ids,
            max_length=768, num_beams=3, early_stopping=True,
        )
    text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    text = text.replace(task_prompt, "").strip()
    lines = [l.strip() for l in text.split("\n") if l.strip()]
    return OCRResult(engine="donut", raw_text=text, lines=lines)


def run_ocr(image: Image.Image, engine: str = "both") -> OCRResult:
    if engine == "trocr":
        return run_trocr(image)
    elif engine == "donut":
        return run_donut(image)
    else:  # both
        donut_res = run_donut(image)
        trocr_res = run_trocr(image)
        combined = f"[Donut]\n{donut_res.raw_text}\n\n[TrOCR]\n{trocr_res.raw_text}"
        return OCRResult("both", combined, donut_res.lines + trocr_res.lines)