DirtyLeague_Bot/ocr_utils.py

"""
OCR utility module using Windows Native OCR (Windows.Media.Ocr via winsdk).
Fast, reliable, offline, and requires no external binaries.
"""

import asyncio
import os
import re
from typing import Optional, Tuple
import cv2
import numpy as np

import winsdk.windows.globalization as glob
import winsdk.windows.graphics.imaging as imaging
import winsdk.windows.media.ocr as win_ocr
import winsdk.windows.storage.streams as streams


class OcrReader:
    def __init__(self, lang_tag: str = "en-US"):
        self.lang_tag = lang_tag
        self._engine = None

    def _get_engine(self):
        if self._engine is None:
            lang = glob.Language(self.lang_tag)
            self._engine = win_ocr.OcrEngine.try_create_from_language(lang)
            if not self._engine:
                self._engine = win_ocr.OcrEngine.try_create_from_user_profile_languages()
        return self._engine

    async def _recognize_async(self, bgr_image: np.ndarray) -> str:
        engine = self._get_engine()
        bgra = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2BGRA)
        success, encoded = cv2.imencode(".png", bgra)
        if not success:
            return ""

        bytes_data = encoded.tobytes()
        stream = streams.InMemoryRandomAccessStream()
        writer = streams.DataWriter(stream)
        writer.write_bytes(bytes_data)
        await writer.store_async()
        await writer.flush_async()
        stream.seek(0)

        decoder = await imaging.BitmapDecoder.create_async(stream)
        software_bitmap = await decoder.get_software_bitmap_async()
        result = await engine.recognize_async(software_bitmap)
        return result.text.strip()

    def read_text(self, bgr_image: np.ndarray) -> str:
        """Synchronously recognizes text from a BGR numpy image."""
        try:
            return asyncio.run(self._recognize_async(bgr_image))
        except Exception:
            # Fallback for environments with active event loop
            loop = asyncio.new_event_loop()
            return loop.run_until_complete(self._recognize_async(bgr_image))

    def read_trophies(self, frame: np.ndarray) -> Optional[Tuple[int, int]]:
        """
        Extracts trophies counter (current, max) from TOWER_LOBBY screen.
        Expected format: '2125/2300' -> returns (2125, 2300).
        """
        # ROI for trophy numbers above FIGHT button
        # y: 540 to 650, x: 3130 to 3500
        fh, fw = frame.shape[:2]
        if fh < 1000 or fw < 2000:
            return None

        trophy_crop = frame[540:650, 3130:3500]
        text = self.read_text(trophy_crop)

        # Clean up text and parse
        match = re.search(r"(\d+)\s*[/\\|]\s*(\d+)", text)
        if match:
            current_val = int(match.group(1))
            max_val = int(match.group(2))
            return current_val, max_val

        return None

    def read_crown_chest(self, frame: np.ndarray) -> Optional[Tuple[int, int]]:
        """
        Extracts crown chest counter (current, max) from TOWER_LOBBY screen.
        Expected format: '3/5' or '5/5' -> returns (current, 5).
        """
        fh, fw = frame.shape[:2]
        if fh < 1000 or fw < 2000:
            return None

        # 1. Fast template check for 5/5 (critical trigger for opening crown chest)
        digit_roi = frame[1390:1465, 790:850]
        tpl_5_path = os.path.join(os.path.dirname(__file__), "assets", "digit_crown_5.png")
        if os.path.exists(tpl_5_path):
            tpl_5 = cv2.imread(tpl_5_path)
            if tpl_5 is not None and digit_roi.shape[0] >= tpl_5.shape[0] and digit_roi.shape[1] >= tpl_5.shape[1]:
                res_5 = cv2.matchTemplate(digit_roi, tpl_5, cv2.TM_CCOEFF_NORMED)
                _, max_v5, _, _ = cv2.minMaxLoc(res_5)
                if max_v5 >= 0.88:
                    return 5, 5

        # 2. Template check for 3/5
        tpl_3_path = os.path.join(os.path.dirname(__file__), "assets", "digit_crown_3.png")
        if os.path.exists(tpl_3_path):
            tpl_3 = cv2.imread(tpl_3_path)
            if tpl_3 is not None and digit_roi.shape[0] >= tpl_3.shape[0] and digit_roi.shape[1] >= tpl_3.shape[1]:
                res_3 = cv2.matchTemplate(digit_roi, tpl_3, cv2.TM_CCOEFF_NORMED)
                _, max_v3, _, _ = cv2.minMaxLoc(res_3)
                if max_v3 >= 0.88:
                    return 3, 5

        # 3. Fallback to Windows native OCR
        badge_crop = frame[1380:1470, 740:1000]
        w = 1000 - 740
        text_part = badge_crop[:, :int(w * 0.65)]
        up = cv2.resize(text_part, (0, 0), fx=3.0, fy=3.0, interpolation=cv2.INTER_CUBIC)
        text = self.read_text(up)

        match = re.search(r"(\d+)\s*[/\\|]\s*(\d+)", text)
        if match:
            current_val = int(match.group(1))
            max_val = int(match.group(2))
            return current_val, max_val

        return None

    def read_timer(self, timer_crop: np.ndarray) -> Optional[int]:
        """
        Reads timer string (e.g. '51:53' or '1:12' or '01:23:45') and converts to total seconds.
        """
        text = self.read_text(timer_crop)
        parts = re.findall(r"\d+", text)
        if len(parts) == 2:
            minutes, seconds = int(parts[0]), int(parts[1])
            return minutes * 60 + seconds
        elif len(parts) == 3:
            hours, minutes, seconds = int(parts[0]), int(parts[1]), int(parts[2])
            return hours * 3600 + minutes * 60 + seconds
        return None


# Global singleton instance for easy access
ocr_reader = OcrReader()


if __name__ == "__main__":
    test_img = cv2.imread("assets/trophy_numbers.png")
    text = ocr_reader.read_text(test_img)
    print("Raw OCR Text:", text)

    full_frame = cv2.imread("assets/trophies_check_snapshot.png")
    trophies = ocr_reader.read_trophies(full_frame)
    print("Parsed Trophies as integers:", trophies)
    if trophies:
        current, maximum = trophies
        print(f"Current: {current} (type: {type(current)}), Max: {maximum} (type: {type(maximum)})")
        print(f"Trophies until goal: {maximum - current}")