import os
from pathlib import Path

from tqdm import tqdm

from recraft_ai import generate_image_v2, sanitize_filename
from replicate_image_models import flux_schnell, nkodeicon_v1, nkodeicon_v3


def run_recraft():
    icon_dir = os.path.expanduser("~/icons")
    os.makedirs(icon_dir, exist_ok=True)
    # style = "nature-simplism"
    with open("prompts.txt", "r") as fp:
        prompts = fp.readlines()
    # completed = 0
    for prompt in tqdm(prompts):
        # if completed == 1:
        #     break
        # file_path = icon_dir + "/" + style + "_" + sanitize_filename(prompt) + ".svg"
        file_path = icon_dir + "/" + sanitize_filename(prompt) + ".svg"
        if os.path.exists(file_path):
            print(f"file exists {file_path}")
            continue
        # generate_image_v3(prompt, Path(file_path), style)
        generate_image_v2(prompt, Path(file_path))
        # completed += 1


def run_flux_schnell():
    icon_dir = os.path.expanduser("~/webp_icons_background")
    os.makedirs(icon_dir, exist_ok=True)
    with open("prompts.txt", "r") as fp:
        prompts = fp.readlines()
    for prompt in tqdm(prompts):
        file_path = icon_dir + "/" + sanitize_filename(prompt) + ".webp"
        if os.path.exists(file_path):
            print(f"file exists {file_path}")
            continue

        flux_schnell(
            f"simple digital art: {prompt}",
            Path(file_path),
        )


def run_nkodeicon_v1():
    icon_dir = os.path.expanduser("~/nkodeiconv1_icons")
    os.makedirs(icon_dir, exist_ok=True)
    with open("prompts.txt", "r") as fp:
        prompts = fp.readlines()
    for prompt in tqdm(prompts):
        file_path = icon_dir + "/" + sanitize_filename(prompt) + ".webp"
        if os.path.exists(file_path):
            print(f"file exists {file_path}")
            continue

        nkodeicon_v1(
            f"{prompt}",
            Path(file_path),
        )


def run_nkodeicon_v3():
    icon_dir = os.path.expanduser("~/nkodeiconv3_icons")
    os.makedirs(icon_dir, exist_ok=True)
    with open("prompts.txt", "r") as fp:
        prompts = fp.readlines()
    for prompt in tqdm(prompts):
        file_path = icon_dir + "/" + sanitize_filename(prompt) + ".webp"
        if os.path.exists(file_path):
            print(f"file exists {file_path}")
            continue

        nkodeicon_v3(
            f"{prompt}",
            Path(file_path),
        )


if __name__ == "__main__":
    run_nkodeicon_v3()