import easyocr
import cv2
import os
from pathlib import Path


class DetectText():
    def __init__(self, logger, threshold=0.0, allowlist=None, text_args=None):
        # Set EasyOCR model storage directory (default: models/easyocr relative to this script)
        default_model_dir = str(Path(__file__).parent / "models" / "easyocr")
        model_storage_directory = os.environ.get('EASYOCR_MODEL_DIR', default_model_dir)

        # This needs to run only once to load the model into memory
        self.reader = easyocr.Reader(['en'], model_storage_directory=model_storage_directory)
        self.threshold = threshold
        self.logger = logger
        self.allowlist = allowlist
        self.text_args = text_args.split(',') if text_args else []

    def detect(self, img):  # expects CV2 image

        if 'threshold' in self.text_args:
            ret, img = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY_INV)

        if 'blur' in self.text_args:
            img = cv2.blur(img, (5, 5))

        if 'grayscale' in self.text_args:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        if 'mag2' in self.text_args:
            mag_ratio = 2.0
        else:
            mag_ratio = 1.0

        output = []
        boxes = []
        # run OCR
        results = self.reader.readtext(img, allowlist=self.allowlist, mag_ratio=mag_ratio)

        for res in results:
            top_left = (int(res[0][0][0]), int(res[0][0][1]))
            bottom_right = (int(res[0][2][0]), int(res[0][2][1]))

            text = res[1]
            confidence = res[2]

            if confidence >= self.threshold:
                output.append((text, confidence))
                boxes.append([top_left, bottom_right])

        return output, boxes