from pytesseract import image_to_string
from PIL import Image, ImageEnhance
from common.read_config import ReadConfig
from guke.common.logger import log


def gen_text_from_picture():
    log.info('开始解析验证码')
    # 打开图片,将图片转换为obj
    image = Image.open(ReadConfig().config.get('SCREENSHOT', 'ELEMENT_SCREENSHOT_PATH'))
    # 识别图中的文字
    verify_code_text = image_to_string(image)
    # 针对平台，去掉最后一个字符“\n”,不同的平台最后一个字符可能不一样
    verify_code_text_no_wrap = verify_code_text.replace('\n', '')
    log.info("当前验证码的文字为：" + verify_code_text_no_wrap)
    # 使用完成后关闭
    image.close()
    # 返回验证码
    return verify_code_text_no_wrap


def gen_text_from_picture_guke():
    log.info('开始解析验证码')
    # 打开图片,将图片转换为obj
    image = Image.open(ReadConfig().config.get('SCREENSHOT', 'ELEMENT_SCREENSHOT_PATH'))
    # 图片2直化
    image = image.convert('L')
    # 灰度化处理
    image = ImageEnhance.Contrast(image)
    image = image.enhance(3)
    image2 = Image.new('RGB', (150, 60), (255, 255, 255))
    image2.paste(image.copy(), (25, 10))
    # 识别图中的文字
    verify_code_text = image_to_string(image2)
    # 针对平台，去掉最后一个字符“\n”,不同的平台最后一个字符可能不一样
    verify_code_text_no_wrap = verify_code_text.replace('\n', '')
    log.info("当前验证码的文字为：" + verify_code_text_no_wrap)
    # 使用完成后关闭
    image.close()
    # 返回验证码
    return verify_code_text_no_wrap


class VerifyCode(object):

    def __init__(self, web_element, driver):
        self.web_element = web_element
        self.driver = driver

    def indentify_verify_image(self):
        # 先刷新验证码
        self.web_element.click()
        # 在首页展开后，对整个页面进行截图,截图保存目录在ini配置文件中定义
        page_shot_dir = ReadConfig().get('SCREENSHOT', 'PAGE_SCREENSHOT_PATH')
        # 截图
        self.driver.save_screenshot(page_shot_dir)
        # 转换图片为Image对象
        page_shot_obj = Image.open(page_shot_dir)
        # 获取验证码元素图片的位置
        location = self.web_element.location
        log.debug("验证码元素图片位置：%s" % location)
        # 获取验证码元素图片的大小
        size = self.web_element.size
        log.debug("验证码元素图片大小：%s" % size)
        # 获取图片的四个坐标点
        left = location['x']
        bottom = location['y']
        right = left + size['width']
        top = bottom + size['height']
        log.debug("图片的四个位置。左:%s右:%s上:%s下:%s" % (left, right, top, bottom))
        # 开始按照位置截图
        crop_image_obj = page_shot_obj.crop((left, bottom, right, top))
        # 将切割好的图片保存，保存的目录为crop_dir，截图保存目录在ini配置文件中定义
        crop_dir = ReadConfig().get('SCREENSHOT', 'ELEMENT_SCREENSHOT_PATH')
        crop_image_obj.save(crop_dir)
        # 识别验证码中的字符
        verify_code_text = gen_text_from_picture()
        log.debug("当前验证码的文字为：%s" % verify_code_text)
        return verify_code_text

    def indentify_verify_image_guke(self):
        # 先刷新验证码
        self.web_element.click()
        # 在首页展开后，对整个页面进行截图,截图保存目录在ini配置文件中定义
        page_shot_dir = ReadConfig().get('SCREENSHOT', 'PAGE_SCREENSHOT_PATH')
        # 截图
        self.driver.save_screenshot(page_shot_dir)
        # 转换图片为Image对象
        page_shot_obj = Image.open(page_shot_dir)
        # 获取验证码元素图片的位置
        location = self.web_element.location
        log.debug("验证码元素图片位置：%s" % location)
        # 获取验证码元素图片的大小
        size = self.web_element.size
        log.debug("验证码元素图片大小：%s" % size)
        # 获取图片的四个坐标点
        left = location['x']
        bottom = location['y']
        right = left + size['width']
        top = bottom + size['height']
        log.debug("图片的四个位置。左:%s右:%s上:%s下:%s" % (left, right, top, bottom))
        # 开始按照位置截图
        crop_image_obj = page_shot_obj.crop((left, bottom, right, top))
        # 将切割好的图片保存，保存的目录为crop_dir，截图保存目录在ini配置文件中定义
        crop_dir = ReadConfig().get('SCREENSHOT', 'ELEMENT_SCREENSHOT_PATH')
        crop_image_obj.save(crop_dir)
        # 识别验证码中的字符
        verify_code_text = gen_text_from_picture_guke()
        log.debug("当前验证码的文字为：%s" % verify_code_text)
        return verify_code_text
