import os

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


def gen_text_from_picture(element_screenshot_path):
    log.info('开始解析验证码')
    # 打开图片,将图片转换为obj
    image = Image.open(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()
        # 在首页展开后，对整个页面进行截图,截图直接保存在根目录autotest_sc_screenshot中
        root_dir = os.path.abspath(os.sep)
        # 截图目录
        shot_dir = os.path.join(root_dir, 'autotest_sc_screenshot')
        # 若该目录不存在则创建该目录
        if not os.path.exists(shot_dir):
            os.mkdir(shot_dir)
        # 设置截图文件名, 页面截图为page.png,元素截图为element.png
        page_screenshot_path = os.path.join(shot_dir, 'page.png')
        element_screenshot_path = os.path.join(shot_dir, 'element.png')
        # 页面截图
        self.driver.save_screenshot(page_screenshot_path)
        # 转换图片为Image对象
        page_shot_obj = Image.open(page_screenshot_path)
        # 获取验证码元素图片的位置
        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_image_obj.save(element_screenshot_path)
        # 识别验证码中的字符
        verify_code_text = gen_text_from_picture(element_screenshot_path)
        log.debug("当前验证码的文字为：%s" % verify_code_text)
        return verify_code_text
