import argparse
import datetime
import os
import uuid
from PIL import Image, ImageFont, ImageDraw
from tqdm import tqdm


def create_img(font_path, text, save_path, cnt):
    fontSize = 25
    # 字体样式，
    font = ImageFont.truetype(font_path, fontSize, encoding="gbk")
    im = Image.new("RGB", (int(font.getbbox(text.rstrip('\n'))[2]), font.getbbox(text)[3]), (255, 255, 255))
    dr = ImageDraw.Draw(im)
    # 文字颜色
    dr.text((0, 0), text, font=font, fill="#000000")
    time = datetime.datetime.now()
    t = str(time.year).zfill(4) + '-' + str(time.month).zfill(2) + '-' + str(time.day).zfill(2)
    if not os.path.exists("%s/%s/" % (save_path, font_path.split('/')[-1][0:-4])):
        os.makedirs("%s/%s/" % (save_path, font_path.split('/')[-1][0:-4]))
    # 图片命名
    img_name = 'kdan_' + t + '_' + str(cnt).zfill(4) + '_' + str(uuid.uuid1())[0:8] + '.jpg'
    save_path = "%s/%s/" % (save_path, font_path.split('/')[-1][0:-4]) + img_name
    # 保存
    im.save(save_path)


def get_all_labels(img_fold):
    # 获取字体文件夹列表
    img_fold_list = os.listdir(img_fold)
    # label_list.txt文件
    fp_label = open(str(img_fold) + '/label_list.txt', 'w')
    cnt = 0
    for img_fold_in in img_fold_list:
        if img_fold_in.endswith('.txt'):
            continue
        # 字体名以及对应id写入label_list.txt
        fp_label.write(str(cnt) + ' ' + str(img_fold_in) + '\n')
        cnt += 1
    fp_label.close()


def work(font_path, text_path, save_path):
    fontPath_list = os.listdir(font_path)
    print('generate images for fonts based on %s' % text_path)
    for fontPath in tqdm(fontPath_list):
        with open(text_path, 'r') as fpp:
            t = fpp.readlines()
        fpp.close()
        cnt = 0
        for text in t:
            create_img(str(font_path) + '/' + str(fontPath), text, save_path, cnt)
            cnt += 1
    get_all_labels(save_path)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--font_dir', type=str, default='./font')
    parser.add_argument('--text_path', type=str, default='./text/eng_text_test.txt')
    parser.add_argument('--save_dir', type=str, default='./images')
    args = parser.parse_args()
    work(args.font_dir, args.text_path, args.save_dir)
