import argparse
import asyncio
import mimetypes
import os
import random
from urllib.parse import urlparse, quote

import aiohttp
from aiofiles import open as aioopen
from playwright.async_api import Page, async_playwright


def get_filename_from_url(url: str, idx: int) -> str:
    parsed = urlparse(url)
    name = os.path.basename(parsed.path)
    if not name:  # 有些URL没有文件名
        name = f"img_{idx}"
    # 如果没有扩展名，尝试补上
    if not os.path.splitext(name)[1]:
        ext = mimetypes.guess_extension(parsed.path.split("?")[0])
        name += ext if ext else ".jpg"
    return name


async def download_image(session: aiohttp.ClientSession, url: str, idx: int,
                         sem: asyncio.Semaphore, save_dir: str, proxy: str | None):
    filename = get_filename_from_url(url, idx)
    filepath = os.path.join(save_dir, filename)

    async with sem:  # 限制并发数
        try:
            async with session.get(url, proxy=proxy) as resp:
                if resp.status == 200:
                    async with aioopen(filepath, "wb") as f:
                        await f.write(await resp.read())
                    print(f"✅ 下载成功: {filepath}")
                else:
                    print(f"❌ 下载失败 {url} 状态码: {resp.status}")
        except Exception as e:
            print(f"⚠️ 下载错误 {url}: {e}")


async def download_all(imgs_url: set[str], save_dir: str, proxy: str | None):
    sem = asyncio.Semaphore(10)  # 同时最多10个下载任务
    async with aiohttp.ClientSession() as session:
        tasks = [
            download_image(session, url, idx, sem, save_dir, proxy)
            for idx, url in enumerate(imgs_url)
        ]
        await asyncio.gather(*tasks)


async def human_move(page: Page, target_x: int, target_y: int):
    """模拟人工鼠标移动: 分步插值 + 随机抖动, 轨迹带弧度"""
    # 获取当前鼠标位置(自己维护, playwright 不提供查询)
    cur_x, cur_y = getattr(human_move, "_pos", (random.randint(100, 800), random.randint(100, 500)))
    steps = random.randint(15, 30)
    # 随机控制点让轨迹带弧度(近似贝塞尔)
    ctrl_x = (cur_x + target_x) / 2 + random.randint(-150, 150)
    ctrl_y = (cur_y + target_y) / 2 + random.randint(-150, 150)
    for i in range(1, steps + 1):
        t = i / steps
        # 二次贝塞尔插值
        x = (1 - t) ** 2 * cur_x + 2 * (1 - t) * t * ctrl_x + t ** 2 * target_x + random.uniform(-2, 2)
        y = (1 - t) ** 2 * cur_y + 2 * (1 - t) * t * ctrl_y + t ** 2 * target_y + random.uniform(-2, 2)
        await page.mouse.move(x, y)
        await page.wait_for_timeout(random.randint(5, 20))  # 毫秒级间隔, 模拟手部移动速度
    human_move._pos = (target_x, target_y)


async def human_scroll(page: Page, distance: int):
    """模拟人工滚动: 把总距离拆成多次小幅滚轮事件, 逐段发出"""
    # 先把鼠标移到页面内一个随机位置再滚
    await human_move(page, random.randint(200, 1000), random.randint(200, 700))
    remaining = distance
    while remaining > 0:
        step = min(random.randint(40, 120), remaining)  # 一次滚轮约 40~120px
        await page.mouse.wheel(0, step)
        remaining -= step
        await page.wait_for_timeout(random.randint(30, 100))  # 滚轮事件间隔


async def scrape(keyword: str, count: int, cdp_url: str) -> set[str]:
    imgs_url: set[str] = set()
    async with async_playwright() as p:
        browser = await p.chromium.connect_over_cdp(cdp_url)
        context = browser.contexts[0]

        page: Page = await context.new_page()
        search_url = f"https://www.pinterest.com/search/pins/?q={quote(keyword)}"
        await page.goto(search_url, wait_until="domcontentloaded")

        no_new_rounds = 0
        while len(imgs_url) < count and no_new_rounds < 10:
            await page.locator('div[role="listitem"]').first.wait_for(state="attached", timeout=50_000)
            before = len(imgs_url)
            imgs = await page.locator('div[role="listitem"]').all()
            print(f"已加载 {len(imgs)} 个元素, 已收集 {len(imgs_url)} 张图片")
            for img in imgs:
                if len(imgs_url) >= count:
                    break
                el = img.locator("img").first
                if not await el.count():
                    continue

                srcset = await el.get_attribute("srcset")
                if not srcset:
                    continue

                candidates = [
                    (s.split()[0], float(s.split()[1][:-1]))  # (url, 倍率)
                    for s in srcset.split(",")
                ]
                max_url = max(candidates, key=lambda x: x[1])[0]
                imgs_url.add(max_url)

            # 模拟人工: 鼠标平滑移动到随机位置 + 分段连续滚动,
            # 平时停 1.5~3 秒, 偶尔长停顿(像在看图)
            await human_scroll(page, random.randint(600, 1400))
            if random.random() < 0.2:
                await page.wait_for_timeout(random.randint(3000, 6000))
            else:
                await page.wait_for_timeout(random.randint(1500, 3000))
            no_new_rounds = no_new_rounds + 1 if len(imgs_url) == before else 0

        # CDP 连接的浏览器不关闭，只断开连接
        await browser.close()

    return imgs_url


async def main():
    parser = argparse.ArgumentParser(description="Pinterest 关键词图片爬取工具 (CDP 连接本地浏览器)")
    parser.add_argument("keyword", help="搜索关键词")
    parser.add_argument("-n", "--count", type=int, default=40, help="爬取图片数量 (默认 40)")
    parser.add_argument("--proxy", default="http://127.0.0.1:7892",
                        help="下载图片使用的代理 (默认 http://127.0.0.1:7892, 传空字符串禁用)")
    parser.add_argument("--cdp", default="ws://127.0.0.1:9222/devtools/browser",
                        help="CDP 接口地址 (默认 ws://127.0.0.1:9222/devtools/browser)")
    parser.add_argument("-o", "--output", default=None, help="保存目录 (默认 output/img/<关键词>)")
    args = parser.parse_args()

    proxy = args.proxy or None
    save_dir = args.output or os.path.join(os.getcwd(), "output", "img", args.keyword)
    os.makedirs(save_dir, exist_ok=True)
    print(f"关键词: {args.keyword} | 数量: {args.count} | CDP: {args.cdp} | 代理: {proxy or '无'}")
    print(f"保存目录: {save_dir}")

    imgs_url = await scrape(args.keyword, args.count, args.cdp)
    await download_all(imgs_url, save_dir, proxy)
    print(f"完成, 共收集 {len(imgs_url)} 张图片")


if __name__ == "__main__":
    asyncio.run(main())