import argparse
import asyncio
import mimetypes
import os
import random
import socket
import sys
import urllib.request
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


def _probe_proxy(url: str) -> bool:
    """实测该地址是否是一个可用的 HTTP/HTTPS 代理（短超时 HEAD 探测）"""
    proxy_handler = urllib.request.ProxyHandler({"http": url, "https": url})
    opener = urllib.request.build_opener(proxy_handler)
    try:
        req = urllib.request.Request(
            "http://www.gstatic.com/generate_204", method="HEAD"
        )
        opener.open(req, timeout=2)
        return True
    except Exception:
        return False


def _read_windows_registry_proxy() -> str | None:
    """直接读取 Windows 系统代理配置（设置 → 网络 → 代理），免去端口扫描猜测。

    来源：注册表 HKCU\\Software\\Microsoft\\Windows\\CurrentVersion\\Internet Settings
      - ProxyEnable == 1 时，ProxyServer 形如 "127.0.0.1:6696" 或 "http=127.0.0.1:8899;https=..."。
    非 Windows 平台或读取失败返回 None。
    """
    try:
        import winreg  # 仅 Windows 可用
    except ImportError:
        return None
    try:
        with winreg.OpenKey(
            winreg.HKEY_CURRENT_USER,
            r"Software\Microsoft\Windows\CurrentVersion\Internet Settings",
        ) as key:
            enabled, _ = winreg.QueryValueEx(key, "ProxyEnable")
            if not enabled:
                return None
            proxy_server, _ = winreg.QueryValueEx(key, "ProxyServer")
        if not proxy_server:
            return None
        # 可能是 "http=127.0.0.1:8899;https=127.0.0.1:8899" 多协议格式，取第一个地址
        first = proxy_server.split(";")[0]
        if "=" in first:
            first = first.split("=", 1)[1]
        return ("http://" + first) if not first.startswith("http") else first
    except Exception:
        return None


def detect_proxy() -> str | None:
    """自动确定本地代理，无需手动填写。

    优先级（确定性强的方式在前，端口扫描兜底在后）：
      1. 系统/用户环境变量 (HTTP_PROXY / HTTPS_PROXY ...)
      2. 直接读 Windows 注册表系统代理 (Internet Settings / ProxyServer) —— 你配的 6696 直接命中
      3. urllib 系统代理 (getproxies，已含注册表/环境变量)
      4. 端口预检 + 实测本地常见代理端口（Clash / v2rayN / Shadowsocks 等）—— 仅作兜底
    探测不到返回 None（调用方将直连）。
    """
    # 1. 环境变量（系统代理 / CI 设置）
    for env in ("HTTPS_PROXY", "https_proxy", "HTTP_PROXY", "http_proxy"):
        val = os.environ.get(env)
        if val:
            return val.rstrip("/")

    # 2. 直接读 Windows 注册表里的系统代理（你手动在系统设置里填的端口，这里精确拿到）
    reg_proxy = _read_windows_registry_proxy()
    if reg_proxy:
        print(f"使用代理: {reg_proxy}（来自系统设置/注册表）")
        return reg_proxy

    # 3. urllib 系统代理兜底（已涵盖注册表/环境变量，跨平台）
    sys_proxy = get_system_proxy()
    if sys_proxy:
        print(f"使用代理: {sys_proxy}（系统代理）")
        return sys_proxy

    # 4. 兜底：实测本地常见代理端口（先快速判断端口是否监听，避免无谓阻塞）
    candidates = [
        # Clash / Clash Verge / Clash for Windows
        "127.0.0.1:7890", "127.0.0.1:7891", "127.0.0.1:7892", "127.0.0.1:7893",
        "127.0.0.1:7894", "127.0.0.1:7878", "127.0.0.1:9090",  # 9090 为 Clash 外部控制(也可能作代理)
        # v2rayN / v2ray-core
        "127.0.0.1:10808", "127.0.0.1:10809", "127.0.0.1:10810", "127.0.0.1:10811",
        "127.0.0.1:10812", "127.0.0.1:10813", "127.0.0.1:10814", "127.0.0.1:10815",
        # Shadowsocks / SSWindows
        "127.0.0.1:1080", "127.0.0.1:1081", "127.0.0.1:1082", "127.0.0.1:1087",
        "127.0.0.1:8388", "127.0.0.1:8389",
        # Surge (mac/iOS 风格，本地也可能开)
        "127.0.0.1:6152", "127.0.0.1:6153",
        # Quantumult / Quantumult X
        "127.0.0.1:6155", "127.0.0.1:6170",
        # 系统代理 / HTTP 调试代理
        "127.0.0.1:8888", "127.0.0.1:8080", "127.0.0.1:8081", "127.0.0.1:8088",
        "127.0.0.1:3128",  # 传统 squid 代理
        # 其他常见：trojan / Brook / Netch / 蓝灯 / 自由门 / Proxifier
        "127.0.0.1:10801", "127.0.0.1:10802", "127.0.0.1:10803", "127.0.0.1:10806",
        "127.0.0.1:8118",  # Privoxy
        "127.0.0.1:10819", "127.0.0.1:2080", "127.0.0.1:1080",
    ]
    for hp in candidates:
        host, port = hp.split(":")
        try:
            with socket.create_connection((host, int(port)), timeout=0.4):
                pass
        except OSError:
            continue  # 端口没开，直接跳过（快速）
        if _probe_proxy(f"http://{hp}"):
            return f"http://{hp}"
    return None


def get_system_proxy() -> str | None:
    """读取操作系统（Windows）设置的代理，作为探测不到本地代理时的兜底。

    Playwright 启动的 Chrome 默认会继承系统代理；但 aiohttp 下载不会，
    因此需要显式取出并传给下载阶段。
    """
    proxies = urllib.request.getproxies()
    val = proxies.get("https") or proxies.get("http")
    return val.rstrip("/") if val else None


def _session_user_data_dir() -> str:
    """返回持久化的 Chrome 用户数据目录。

    该目录会一直保留（已加入 .gitignore），首次手动登录 Pinterest 后，
    登录态（cookies 等）自动缓存在这里，后续运行直接复用，不再需要登录。
    源码运行时：pinterest_scraper/.chrome_session（项目内）；
    打包 exe 运行时：exe 同目录 pinterest_scraper/.chrome_session（持久，
    打包时把当前登录态复制过去即可复用，不随临时解压目录消失）。
    """
    if getattr(sys, "frozen", False):
        base = os.path.dirname(os.path.abspath(sys.executable))
        ud = os.path.join(base, "pinterest_scraper", ".chrome_session", "User Data")
    else:
        base = os.path.dirname(os.path.abspath(__file__))
        ud = os.path.join(base, ".chrome_session", "User Data")
    os.makedirs(ud, exist_ok=True)
    return ud


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 _ensure_logged_in(page: Page, search_url: str) -> bool:
    """独立登录检查方法：在已跳到搜索页的前提下确认 Pinterest 已登录。

    判定策略（避免误判）：
    - 以「未登录标志」为准：页面上一旦出现 "Log in" / "Sign up" 按钮，
      才认定未登录；否则默认已登录（不依赖可能不匹配的已登录选择器）。
    - Pinterest 是 SPA，goto 后需等待渲染，否则瞬间误判未登录。
    - 已登录 → 立即返回 True，走直路。
    - 未登录（被弹回登录墙）→ 回退首页提示手动登录一次，登录成功后返回 True。
    - 一直未登录（用户关窗口/放弃）→ 返回 False。
    """
    async def _has_login_wall() -> bool:
        """检测是否存在未登录标志（Log in / Sign up 按钮）。"""
        try:
            if await page.get_by_role("button", name="Log in").count():
                return True
        except Exception:
            pass
        try:
            if await page.get_by_role("button", name="Sign up").count():
                return True
        except Exception:
            pass
        # 兜底：URL 被重定向到 /login 也是未登录的强信号
        try:
            if "/login" in page.url:
                return True
        except Exception:
            pass
        return False

    # 已经在搜索页了。等待 SPA 渲染，避免刚加载就被误判。
    # 给一点时间让导航/按钮渲染出来（最多等 8 秒，出现登录墙或超时即停）。
    try:
        await page.wait_for_load_state("networkidle", timeout=8000)
    except Exception:
        pass  # 网络一直不 idle 也不要卡死，继续判断

    # 已登录：页面上找不到登录墙 → 直接走直路
    if not await _has_login_wall():
        print("✅ 已检测到登录态，直接开始搜索")
        return True

    # 未登录（搜索页被弹回登录墙）：回退首页，提示手动登录一次，成功后立即继续
    print("⚠️ 当前未登录（搜索页被拦截）。请在弹出的浏览器窗口中手动登录，登录成功后将自动继续……")
    try:
        await page.goto("https://www.pinterest.com/", wait_until="domcontentloaded")
        # 登录成功后登录墙消失（Log in/Sign up 按钮不再存在）即视为登录
        await page.wait_for_function(
            """() => {
                const btns = [...document.querySelectorAll('button')];
                const hasLogin = btns.some(b => /log\\s*in/i.test(b.textContent || ''));
                const hasSignup = btns.some(b => /sign\\s*up/i.test(b.textContent || ''));
                return !hasLogin && !hasSignup;
            }""",
            timeout=0,  # 0 = 一直等到出现为止，不超时
        )
        print("✅ 登录成功，继续搜索")
        return True
    except Exception:
        # 用户关掉页面 / 主动放弃
        print("❌ 未检测到登录（页面已关闭或放弃登录）。请先登录后再运行脚本。")
        return False


async def scrape(keyword: str, count: int, headless: bool = False,
                 proxy: str | None = None,
                 cdp_url: str | None = None) -> set[str]:
    imgs_url: set[str] = set()
    async with async_playwright() as p:
        if cdp_url:
            # 高级模式：直接连用户已开好的调试版主浏览器（如 open_chrome_debug.bat）
            browser = await p.chromium.connect_over_cdp(cdp_url)
            context = browser.contexts[0] if browser.contexts else await browser.new_context()
            print(f"已接管本机调试 Chrome: {cdp_url}")
        else:
            # 默认模式：用项目内持久化的 Chrome 用户数据目录启动一个可见窗口。
            # 首次运行请在弹出的窗口里登录 Pinterest；登录态会缓存在
            # .chrome_session 目录中，之后运行自动复用，无需再次登录。
            # 复用系统已装的 Chrome（channel="chrome"），无需下载浏览器内核。
            # 注意：Playwright 设置用户数据目录必须用 launch_persistent_context，
            # 它返回的是 context（而非 browser），且该 context 已带登录态。
            launch_kwargs = {
                "headless": headless,
                "channel": "chrome",
                "user_data_dir": _session_user_data_dir(),
                "args": [
                    "--disable-blink-features=AutomationControlled",
                    "--no-first-run",
                    "--no-default-browser-check",
                ],
            }
            if proxy:
                launch_kwargs["proxy"] = {"server": proxy}
            context = await p.chromium.launch_persistent_context(**launch_kwargs)
            browser = context.browser
            print("已启动本地 Chrome 窗口（可见，登录态缓存在 .chrome_session）")

        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")

        # 独立登录检查：已登录直接开始；未登录才回退首页等手动登录
        logged_in = await _ensure_logged_in(page, search_url)
        if not logged_in:
            # 关闭浏览器（持久化目录已保存任何已有状态），中止本次爬取
            if cdp_url:
                await browser.close()
            else:
                await context.close()
            return set()

        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

        # 关闭浏览器（持久化目录中的登录态已自动保存，下次运行直接复用）
        if cdp_url:
            await browser.close()  # CDP 接管模式：只断开，不关主浏览器
        else:
            await context.close()  # 持久化 context 模式：关闭即保存登录态

    return imgs_url


async def main():
    parser = argparse.ArgumentParser(description="Pinterest 关键词图片爬取工具 (本地 Chrome + 持久化登录缓存)")
    parser.add_argument("keyword", help="搜索关键词")
    parser.add_argument("-n", "--count", type=int, default=40, help="爬取图片数量 (默认 40)")
    parser.add_argument("--proxy", default=None,
                        help="下载/浏览使用的代理 (默认自动探测本地端口; 传空字符串禁用)")
    parser.add_argument("--cdp", default=None,
                        help="可选：指定 CDP 地址以接管你已用调试模式打开的主浏览器 "
                             "(见 open_chrome_debug.bat)。不填则默认用项目内持久化缓存目录启动 Chrome，"
                             "首次登录后自动缓存登录态")
    parser.add_argument("--headless", action="store_true",
                        help="自动启动 Chrome 时采用无头模式 (默认显示浏览器窗口)")
    parser.add_argument("-o", "--output", default=None, help="保存目录 (默认 output/img/<关键词>)")
    args = parser.parse_args()

    # 代理解析（命令行 > 本地端口探测 > 系统代理 > 直连）；空字符串明确禁用
    proxy = args.proxy
    if proxy is None:
        proxy = detect_proxy() or get_system_proxy()
        if proxy:
            print(f"使用代理: {proxy}（本地探测/系统代理）")
        else:
            print("未检测到代理, 将直连 (如无法访问 Pinterest 请手动指定 --proxy)")
    proxy = 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} | 模式: {'接管主浏览器' if args.cdp else '本地窗口(缓存登录态)'} | 代理: {proxy or '无'}")
    print(f"保存目录: {save_dir}")

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


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