import time

from airtest.core.cv import try_log_screen
from airtest.core.error import TargetNotFoundError
from airtest.core.helper import logwrap, G
from airtest.core.settings import Settings as ST

from zafkiel.logger import logger
from zafkiel.ocr.ocr import Ocr
from zafkiel.utils import is_color_similar, crop


@logwrap
def loop_find(v, timeout=ST.FIND_TIMEOUT, threshold=None, interval=0.5, interval_func=None, ocr_mode=0, cls=Ocr):
    """
    Search for image template in the screen until timeout
    Add OCR and color similarity search to airtest.cv.loop_find()

    Args:
        v: image template to be found in screenshot
        timeout: time interval how long to look for the image template
        threshold: default is None
        interval: sleep interval before next attempt to find the image template
        interval_func: function that is executed after unsuccessful attempt to find the image template
        ocr_mode: Ocr match rules, one of `OCR_EQUAL`, `OCR_CONTAINS`, `OCR_SIMILAR`.
        cls: "Ocr" class or its subclass

    Raises:
        TargetNotFoundError: when image template is not found in screenshot

    Returns:
        TargetNotFoundError if image template not found, otherwise returns the position where the image template has
        been found in screenshot
    """
    start_time = time.time()
    while True:
        screen = G.DEVICE.snapshot(filename=None, quality=ST.SNAPSHOT_QUALITY)

        if screen is None:
            logger.warning("Screen is None, may be locked")
        else:
            if threshold:
                v.threshold = threshold

            if v.keyword is not None:
                ocr = cls(v)
                if ocr.ocr_match_keyword(screen, ocr.button.keyword, mode=ocr_mode):
                    match_pos = int((v.area[0] + v.area[2]) / 2), int((v.area[1] + v.area[3]) / 2)
                    try_log_screen(screen)
                    return match_pos
            elif not v.rgb or is_color_similar(v.image, crop(screen, v.area)):
                match_pos = v.match_in(screen)
                if match_pos:
                    try_log_screen(screen)
                    return match_pos

        if interval_func is not None:
            interval_func()

        if (time.time() - start_time) > timeout:
            try_log_screen(screen)
            raise TargetNotFoundError(f'Picture {v} not found on screen')
        else:
            time.sleep(interval)
