from auto_media_publisher.utils.utils_text2image_pollinations import Text2ImageGenerator
from auto_media_publisher.utils.utils_file import load_json
import os
import random
import datetime
import sys
import time

def main():

    count = 100

    if len(sys.argv) >1:
        try:
            count = int(sys.argv[1])
            if count >100 or count < 1:
                count =100
        except:
            pass

    generator = Text2ImageGenerator()

    # 待生成的中文提示词（必填）
    prompts = load_json("./scripts/office.json")

    output_dir = "./scripts/output"
    os.makedirs(output_dir,exist_ok=True)
    # 生成并保存图像
    try:
        for i in range(count):
            prompt = random.choice(prompts)
            now = datetime.datetime.now()
            filename = os.path.join(output_dir,now.strftime("%Y%m%d%H%M%S") + ".png")

            saved_path = generator.generate_image_without_trans(prompt, filename=filename)
            print(f"[√] 图片成功保存至：{saved_path}")
            time.sleep(random.randint(60,80))

    except Exception as e:
        print(f"[×] 生成失败：{e}")

if __name__ == "__main__":
    main()