import asyncio
import aiohttp
import os
import io
import re
import sys
from PIL import Image

# =========================
# Парсинг ввода: URL или fetch(...) из DevTools
# =========================

def parse_fetch_snippet(raw_text):
    """
    Принимает либо голый URL, либо скопированный из DevTools (правый клик на
    запросе -> Copy -> Copy as fetch / Copy as Node.js fetch) блок вида:

    fetch("https://...", {
      "headers": { "accept": "...", ... },
      "referrer": "https://...",
      "method": "GET",
      ...
    });

    Возвращает (url_template, headers_dict).
    Если это просто ссылка без fetch(...) - вернёт (url, {}).
    """
    raw_text = raw_text.strip()

    # Если это просто голая ссылка (нет слова fetch(...) в начале) -
    # считаем её обычным URL, даже если в ней есть "{}" как плейсхолдер номера.
    if not raw_text.lower().startswith("fetch("):
        return raw_text, {}

    # Достаём первый аргумент fetch("URL", ...)
    url_match = re.search(r'fetch\(\s*["\'](.+?)["\']\s*(?:,|\))', raw_text, re.DOTALL)
    if not url_match:
        # на случай, если просто вставили URL в кавычках
        url_match = re.search(r'["\'](https?://[^"\']+)["\']', raw_text)
    if not url_match:
        raise ValueError("Не удалось найти URL в переданном тексте")

    url = url_match.group(1)

    headers = {}

    # Вырезаем блок "headers": { ... } (учитываем один уровень вложенности)
    headers_match = re.search(r'["\']headers["\']\s*:\s*\{(.*?)\}\s*,?', raw_text, re.DOTALL)
    if headers_match:
        headers_block = headers_match.group(1)
        # Ищем пары "key": "value"
        pair_re = re.compile(r'["\']([^"\']+)["\']\s*:\s*["\']((?:[^"\'\\]|\\.)*)["\']')
        for key, value in pair_re.findall(headers_block):
            value = value.replace('\\"', '"').replace("\\'", "'")
            headers[key] = value

    # Отдельно вытаскиваем referrer, если он есть вне headers
    referrer_match = re.search(r'["\']referrer["\']\s*:\s*["\']([^"\']*)["\']', raw_text)
    if referrer_match and referrer_match.group(1):
        headers.setdefault("Referer", referrer_match.group(1))

    # DevTools не включает User-Agent в fetch-снипет (браузер подставляет
    # его сам), но многим сайтам (например pixiv) он нужен, иначе 403.
    if not any(k.lower() == "user-agent" for k in headers):
        headers["User-Agent"] = (
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
            "(KHTML, like Gecko) Chrome/148.0.0.0 Safari/537.36"
        )

    return url, headers


def build_url_template(url):
    """
    Заменяет число прямо перед расширением файла на {}, чтобы получить шаблон.
    Пример: https://i.pximg.net/.../145654493_p0.jpg -> оставляем как есть,
    если пользователь сам не поставил {} - подставлять шаблон некуда, поэтому
    просто возвращаем URL как шаблон (он либо уже содержит {}, либо будет
    использован без подстановки номера, если ссылка не подходит под перебор).
    """
    return url


# =========================
# Настройки (ввод пользователя)
# =========================

print("Вставьте ссылку или fetch(...) из DevTools (F12 -> Network -> Copy -> Copy as fetch / Copy as Node.js fetch).")
print("В URL используйте {} на месте числа, которое нужно перебирать, например:")
print("  https://site.com/img/{}.webp")
print("  fetch(\"https://i.pximg.net/img-original/img/2026/06/06/02/29/52/{}_p0.jpg\", {{ ...headers... }})")
print()

raw_input_text = []
print("Вставьте текст (можно многострочный для fetch), затем пустую строку для завершения:")
while True:
    try:
        line = input()
    except EOFError:
        break
    if line == "" and raw_input_text:
        break
    raw_input_text.append(line)

raw_text = "\n".join(raw_input_text).strip()

url_template, headers = parse_fetch_snippet(raw_text)

if "{}" not in url_template:
    print("\n[!] В ссылке нет '{}' - число для перебора подставляться не будет,")
    print("    будет скачан только один файл по указанному URL.\n")

if headers:
    print(f"\nНайдено заголовков: {len(headers)}")
    for k in headers:
        print(f"  - {k}")
    print()

start_input = input("Начальный номер (Enter = 1): ").strip()
start = int(start_input) if start_input else 1

end_input = input("Конечный номер (Enter = автоопределение): ").strip()
end = int(end_input) if end_input else None

folder_name = input("Название папки (Enter = images): ").strip()
if not folder_name:
    folder_name = "images"

os.makedirs(folder_name, exist_ok=True)

CONCURRENT_DOWNLOADS = 20

# =========================
# Конвертация в JPEG
# =========================

def convert_to_jpg(raw_bytes):
    img = Image.open(io.BytesIO(raw_bytes))

    # Сохраняем прозрачность если есть
    # Сохраняем прозрачность WebP на белом фоне
    if img.mode in ("RGBA", "PA", "P"):
        bg = Image.new("RGB", img.size, (255, 255, 255))
        if img.mode == "P":
            img = img.convert("RGBA")
        if img.mode in ("RGBA", "PA"):
            bg.paste(img, mask=img.split()[-1])
        img = bg
    elif img.mode != "RGB":
        img = img.convert("RGB")

    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=95, optimize=True)
    return buf.getvalue()

# =========================
# Асинхронное скачивание
# =========================

async def download_image(session, number):
    url = url_template.format(number) if "{}" in url_template else url_template

    try:
        async with session.get(url, headers=headers) as response:
            if response.status == 200:
                raw = await response.read()

                try:
                    jpg_bytes = convert_to_jpg(raw)
                    filename = os.path.join(folder_name, f"{number}.jpg")
                    with open(filename, "wb") as f:
                        f.write(jpg_bytes)
                    print(f"[OK] {number} -> {number}.jpg")
                except Exception as e:
                    print(f"[CONV ERROR] {number} -> {e}, сохраняю оригинал")
                    ext = url.rsplit(".", 1)[-1].split("?")[0] or "bin"
                    filename = os.path.join(folder_name, f"{number}.{ext}")
                    with open(filename, "wb") as f:
                        f.write(raw)

                return True

            else:
                print(f"[MISS] {number} -> HTTP {response.status}")
                return False

    except Exception as e:
        print(f"[ERROR] {number} -> {e}")
        return False

# =========================
# Автоопределение конца
# =========================

async def find_last_image(session, start_number):
    print("Автоопределение последней картинки...")
    current = start_number
    while True:
        url = url_template.format(current)
        try:
            async with session.get(url, headers=headers) as response:
                if response.status != 200:
                    return current - 1
        except:
            return current - 1
        current += 1

# =========================
# Основная логика
# =========================

async def main():
    connector = aiohttp.TCPConnector(limit=CONCURRENT_DOWNLOADS)

    async with aiohttp.ClientSession(connector=connector, trust_env=True) as session:
        global end

        if "{}" not in url_template:
            # Одиночная ссылка без перебора номеров
            print(f"\nСкачиваю одиночный файл: {url_template}\n")
            await download_image(session, start)
            print("\nГотово.")
            return

        if end is None:
            end = await find_last_image(session, start)

        print(f"\nДиапазон: {start} -> {end}\n")

        tasks = [
            asyncio.create_task(download_image(session, number))
            for number in range(start, end + 1)
        ]

        await asyncio.gather(*tasks)

    print("\nГотово.")


asyncio.run(main())