import os
import time
import requests


# 是否启用 SOCKS5 代理
USE_PROXY = True
# SOCKS5 代理地址
SOCKS5_PROXY = "192.168.2.205:1080"

# 代理探测重试配置
MAX_RETRIES = 3  # 最大重试次数
RETRY_DELAY = 2  # 每次失败后等待秒数

# GitHub 仓库信息
REPO = "ggml-org/llama.cpp"

# 构造代理字典 (使用 socks5h 以便在远程进行 DNS 解析)
PROXIES = (
    {
        "http": f"socks5h://{SOCKS5_PROXY}",
        "https": f"socks5h://{SOCKS5_PROXY}",
    }
    if USE_PROXY
    else None
)


def check_proxy():
    """探测 SOCKS5 代理是否正常，包含重试机制"""
    if not USE_PROXY:
        return True

    print(f"SOCKS5 {SOCKS5_PROXY}")

    for i in range(1, MAX_RETRIES + 1):
        try:
            # 探测目标选择 github.com，确保代理能上外网
            # 使用 HEAD 请求，设置 5 秒超时
            response = requests.head("https://github.com", proxies=PROXIES, timeout=3)

            if response.status_code < 400:
                print(f"第 {i} 次尝试，代理连接成功！")
                return True
            else:
                print(f"第 {i} 次尝试，状态码异常 ({response.status_code})")
        except Exception as e:
            print(f"第 {i} 次尝试，失败 (错误: {e})")

        # 如果不是最后一次尝试，则等待后再试
        if i < MAX_RETRIES:
            print(f"等待 {RETRY_DELAY} 秒后重试...")
            time.sleep(RETRY_DELAY)

    return False


def get_latest_release():
    """获取最新版本的 tag 名称"""
    api_url = f"https://api.github.com/repos/{REPO}/releases/latest"

    print(f"\n正在从 GitHub API 获取 {REPO} 的最新版本...")
    try:
        response = requests.get(api_url, proxies=PROXIES, timeout=15)
        response.raise_for_status()
        data = response.json()
        return data["tag_name"]
    except Exception as e:
        print(f"获取版本信息失败: {e}")
        return None


def download_file(url, filename):
    """下载文件并显示进度"""
    print(f"\n准备下载: {filename}")
    try:
        with requests.get(url, proxies=PROXIES, stream=True, timeout=30) as r:
            r.raise_for_status()
            total_size = int(r.headers.get("content-length", 0))

            with open(filename, "wb") as f:
                chunk_size = 8192
                downloaded = 0
                for chunk in r.iter_content(chunk_size=chunk_size):
                    f.write(chunk)
                    downloaded += len(chunk)
                    if total_size > 0:
                        done = int(50 * downloaded / total_size)
                        print(
                            f"\r进度: [{'█' * done}{'.' * (50 - done)}] {downloaded/(1024*1024):.2f}MB / {total_size/(1024*1024):.2f}MB",
                            end="",
                        )
                print("\n下载完成！")
    except Exception as e:
        print(f"\n下载失败: {e}")


def main():
    # 0. 代理预热/探测
    if USE_PROXY:
        if not check_proxy():
            print("\n[错误] 经过多次尝试，代理服务器仍无法正常工作。")
            print(
                "检查确认：1. SOCKS5 服务已启动；2. 地址和端口正确；3. 网络路径通畅。"
            )
            return

    # 1. 获取最新版本号
    tag = get_latest_release()
    if not tag:
        return

    print(f"检测到最新版本: {tag}")

    # 2. 定义要下载的文件列表
    files_to_download = [
        "cudart-llama-bin-win-cuda-13.1-x64.zip",
        f"llama-{tag}-bin-win-cuda-13.1-x64.zip",
    ]

    # 3. 执行下载
    base_url = f"https://github.com/{REPO}/releases/download/{tag}"

    for file_name in files_to_download:
        download_url = f"{base_url}/{file_name}"
        if os.path.exists(file_name):
            print(f"\n文件 {file_name} 已存在，跳过。")
            continue
        download_file(download_url, file_name)


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"\n运行出错: {e}")

    input(f"\n任务结束，按回车键退出...")