import random
import logging
import base64
from os import urandom
from io import BytesIO
from PIL import ImageFilter
from captcha.image import ImageCaptcha, random_color

__author__ = 'Woodstock'

__doc__ = """
验证码图片生成
"""

logger = logging.getLogger(__name__)


class CaptchaCode(object):
    BASES = '02345689'  # 1acdefhjkmnprtuvwxyACDEFGHJKLMNPQRTUVWXY'  # 去除部分不易识别的字母

    @staticmethod
    def generate(code_len=4):
        """
        验证码长度
        :param code_len:
        :return:
        """
        code = urandom(code_len)  # 验证码长度
        n = len(CaptchaCode.BASES)
        return "".join([CaptchaCode.BASES[int(c) % n - 1] for c in code])


class CustomCaptcha(ImageCaptcha):
    def __init__(self, width=160, height=60, fonts=None, font_sizes=None):
        super(CustomCaptcha, self).__init__(width, height, fonts,
                                            font_sizes=self._scale_font(font_sizes, width, height))

    def _scale_font(self, font_sizes, width, height):
        """
        :param font_sizes:
        :param width:
        :param height:
        :return:
        """
        if not font_sizes:
            r = width / 320 + height / 120
            font_sizes = set([int(fs * r) for fs in (42, 50, 56)])
        logger.debug(str(font_sizes))
        return font_sizes

    def generate_image_custom(self, chars, dots_count=30, dots_size=3, curve=1):
        """
        :param chars:
        :param dots_count:
        :param dots_size:
        :param curve:
        :return:
        """
        background = random_color(238, 255)
        color = random_color(0, 200, random.randint(220, 255))
        im = self.create_captcha_image(chars, color, background)
        if dots_count > 0 and dots_size > 0:
            self.create_noise_dots(im, color, dots_size, dots_count)

        if curve:
            self.create_noise_curve(im, color)
        im = im.filter(ImageFilter.SMOOTH)
        buf = BytesIO()
        im.save(buf, format='PNG')
        buf.seek(0)
        return buf

    def generate_image_base64_str(self, chars, dots_count=30, dots_size=3, curve=1):
        """
        生成图片验证码二进制字符串
        :param chars:
        :param dots_count:
        :param dots_size:
        :param curve:
        :return:
        """
        buf = self.generate_image_custom(chars, dots_count, dots_size, curve)
        byte_data = buf.getvalue()
        base64_str = str(base64.b64encode(byte_data), 'utf-8')
        return base64_str


if __name__ == '__main__':
    captcha_code = CaptchaCode.generate()
    print(captcha_code)
    base64_str = CustomCaptcha().generate_image_base64_str(captcha_code)
    print(base64_str)
    print(len(base64_str))