import captcha_settings
import os
import random

from captcha.image import ImageCaptcha
from PIL import Image

from tqdm import trange

def random_captcha_text(char_set=captcha_settings.NUMBER + captcha_settings.ALPHABET, captcha_size=4):
    captcha_text = []
    for i in range(captcha_size):
        c = random.choice(char_set)
        captcha_text.append(c)
    return "".join(captcha_text)

def gen_captcha_text_and_image():
    image = ImageCaptcha()
    captcha_text = random_captcha_text()
    captcha_image = Image.open(image.generate(captcha_text))
    return captcha_text, captcha_image

if __name__ == "__main__":
    dataset_type = input("请输入数据集类型（1 - train / 2 - test / 3 - predict）：")
    dataset_len = input("请输入数据集长度：")

    paths = [captcha_settings.TRAIN_DATASET_PATH, captcha_settings.TEST_DATASET_PATH, captcha_settings.PREDICT_DATASET_PATH]

    dataset_type = int(dataset_type)
    count = int(dataset_len)
    path = (
        paths[dataset_type - 1]
    )
    if not os.path.exists(path):
        os.makedirs(path)
    for i in trange(count):
        text, image = gen_captcha_text_and_image()
        filename = f"{str(i).zfill(5)}_{text}.png"
        image.save(path + os.path.sep + filename)