import os
import requests
from concurrent.futures import ThreadPoolExecutor
from tqdm import tqdm


local = 0
endpoint = "http://127.0.0.1:7860" if local else "https://q6-p.hf.space"
img_base = 'https://i.pximg.net/img-original/img/'

os.chdir(os.path.dirname(os.path.abspath(__file__)))

valid = [f for f in os.listdir() if f.endswith(".txt")]
for idx, file in enumerate(valid):
    print(f"{idx + 1}: {file}")

inputs = input("Enter the index of the file: ").split()
indexes = []
for inp in inputs:
    if "-" in inp:
        start, end = map(int, inp.split("-"))
        indexes.extend(range(start - 1, end))
    elif inp.isdigit():
        indexes.append(int(inp) - 1)

def download_image(args):
    idx, url, group_name = args
    r = requests.get(url, headers={"Referer": "https://www.pixiv.net/",
                                   "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:126.0) Gecko/20100101 Firefox/126.0"})
    with open(f"images/{group_name}/{idx}.png", "wb") as f:
        f.write(r.content)

for i in indexes:
    group_name = valid[i].rsplit(".", 1)[0]
    os.makedirs(f"images/{group_name}", exist_ok=True)
    with open(valid[i], "r") as f:
        post_ids = [line.strip() for line in f if line.strip()]
    if not post_ids:
        continue
    params = [("post_ids", pid) for pid in post_ids] + [("only_first", "1")]
    data = requests.get(f"{endpoint}/allimages", params=params).json()
    image_urls = data.get("image_urls", [])
    if not image_urls:
        continue
    to_download = [(idx, url, group_name) for idx, url in enumerate(image_urls)]
    with ThreadPoolExecutor(max_workers=20) as executor:
        list(tqdm(executor.map(download_image, to_download), total=len(to_download), desc=f"Downloading {group_name}"))
    print(f"Saved {len(image_urls)} images to images/{group_name}")