import argparse
import json
from datetime import datetime
from playwright.sync_api import sync_playwright


def fetch_and_filter_templates_via_browser(top_n=5, sort_by="downloads"):
    url = "https://cloud.vast.ai/templates/"

    print("正在启动浏览器并加载 Vast.ai 模板页面 (请稍候)...")

    with sync_playwright() as p:
        # 启动无头浏览器
        browser = p.chromium.launch(headless=True)
        page = browser.new_page()

        # 拦截或直接等待页面加载
        try:
            page.goto(url, wait_until="networkidle", timeout=30000)
        except Exception as e:
            print(f"页面加载超时或失败: {e}")
            browser.close()
            return

        print("页面加载完成，正在提取模板数据...")

        # 核心：从前端的全局变量或 DOM 中提取渲染好的数据
        # Vast.ai 的 React/Vue 状态通常挂载在 window 上，或者我们可以直接抓取渲染出来的卡片
        # 这里我们直接抓取页面上包含镜像名称的容器元素
        try:
            # 等待镜像卡片渲染完成（根据实际网页结构调整选择器）
            page.wait_for_selector("div", timeout=5000)

            # 稳妥起见：我们在页面上下文中执行一段 JS，把前端缓存的数据捞出来
            # 如果撈不到，我们就通过解析网络请求（Network Interception）来截获数据
            # 下面使用拦截网络请求的方式，这是最 100% 成功的方法：
        except Exception as e:
            print(f"解析页面失败: {e}")

        browser.close()


# =====================================================================
# 备用高成功率方案：利用 Playwright 监听真实的 API 请求
# =====================================================================
def fetch_by_intercepting_network(top_n=5, sort_by="downloads"):
    url = "https://cloud.vast.ai/templates/"
    templates_data = []

    with sync_playwright() as p:
        browser = p.chromium.launch(headless=True)
        page = browser.new_page()

        # 监听所有的网络响应
        def handle_response(response):
            # 寻找 URL 中包含 templates 的真正请求
            if "templates" in response.url and response.status == 200:
                try:
                    json_data = response.json()
                    # 如果返回的是字典且含有 templates 键
                    if (
                        isinstance(json_data, dict)
                        and "templates" in json_data
                    ):
                        nonlocal templates_data
                        templates_data = json_data["templates"]
                except Exception:
                    pass

        page.on("response", handle_response)

        print("正在通过浏览器截获真实的 Templates API 请求...")
        try:
            page.goto(url, wait_until="networkidle", timeout=20000)
        except Exception:
            # networkidle 有时会因为某些埋点请求未结束而超时，如果已经拿到数据了可以忽略
            pass

        browser.close()

    if not templates_data:
        print(
            "未能截获到模板数据，请检查网络是否能正常访问 https://cloud.vast.ai/templates/"
        )
        return

    # 数据清洗与排序
    parsed = []
    for t in templates_data:
        image_name = t.get("image", "Unknown")
        downloads = t.get("downloads", 0)
        # 兼容处理时间
        updated_ts = t.get("updated", 0)
        try:
            # 有些是秒级时间戳，有些是毫秒级，做个判断
            ts = updated_ts / 1000 if updated_ts > 1e11 else updated_ts
            updated_time = datetime.fromtimestamp(ts).strftime(
                "%Y-%m-%d %H:%M:%S"
            )
        except Exception:
            updated_time = "Unknown"

        parsed.append(
            {
                "image": image_name,
                "downloads": downloads,
                "updated_ts": updated_ts,
                "updated_time": updated_time,
                "tag": image_name.split(":")[-1]
                if ":" in image_name
                else "latest",
            }
        )

    # 排序逻辑
    if sort_by == "downloads":
        filtered = sorted(parsed, key=lambda x: x["downloads"], reverse=True)
        label = "下载量最多"
    else:
        filtered = sorted(parsed, key=lambda x: x["updated_ts"], reverse=True)
        label = "最近更新"

    # 输出结果
    result = filtered[:top_n]
    print(f"\n=== 成功截获数据！{label} 的 Top {top_n} 镜像 ===")
    print(
        f"{'排名':<4} | {'Docker 镜像名称':<50} | {'Tag':<15} | {'下载量':<10} | {'更新时间':<20}"
    )
    print("-" * 110)
    for idx, item in enumerate(result, 1):
        print(
            f"{idx:<4} | {item['image'].split(':')[0]:<50} | {item['tag']:<15} | {item['downloads']:<10} | {item['updated_time']:<20}"
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Vast.ai 模板动态爬虫")
    parser.add_argument(
        "-n", type=int, default=5, help="Top N 数量 (默认: 5)"
    )
    parser.add_argument(
        "--sort",
        choices=["downloads", "updated"],
        default="downloads",
        help="排序: downloads 或 updated",
    )
    args = parser.parse_args()

    fetch_by_intercepting_network(top_n=args.n, sort_by=args.sort)