import json
import os
import shutil
from concurrent.futures import ThreadPoolExecutor
from threading import Lock

import requests


# 获取文件大小
def length(url_list):
    def getlenhead(single_url):
        response = requests.head(single_url)
        while "Location" in response.headers:
            response = requests.head(response.headers.get("Location"))
        file_size = response.headers.get("Content-Length")
        if file_size is not None:
            return int(file_size)
        else:
            return False

    for url in url_list:
        try:
            single_file_length = getlenhead(url)
            if single_file_length:
                return single_file_length
        except Exception as e:
            print(f"{url} 无法连接")
            print(e)


# 定义Download类在初始化时保存几个参数
class Downloader:
    # 初始化类
    def __init__(self, urlist, chunksize, max_conn, use_proxies=None):
        self.urlist = urlist  # 镜像url列表
        self.chunksize = chunksize  # 分片大小
        self.max_conn = max_conn  # 单个url最大连接数
        self.proxies = use_proxies  # 代理服务器
        self.lock = Lock()
        self.chunk_status = []  # 状态列表
        self.failed_requests = {
            url: {"success": 0, "fail": 0} for url in urlist
        }  # 记录每个 URL 的失败次数和成功次数
        self.listhash = hex(hash(tuple(urlist)))  # 计算urlist的hash

    def download_chunk(self, url, chunk_id, total_size):
        start = chunk_id * self.chunksize
        end = min(start + self.chunksize - 1, total_size - 1)
        headers = {"Range": f"bytes={start}-{end}"}
        filename = f"temp/{self.listhash}/{chunk_id}"
        if self.chunk_status[chunk_id] != 2:
            try:
                response = requests.get(
                    url, headers=headers, proxies=self.proxies, timeout=5
                )
                if response.status_code in (
                    301,
                    302,
                    303,
                    307,
                    308,
                ):  # 处理HTTP 3xx 重定向问题，继续发送原来的header（range）
                    redirect_url = response.headers["Location"]
                    response = requests.get(
                        redirect_url, headers=headers, timeout=3, proxies=self.proxies
                    )

                if response.status_code == 206:  # 206状态码
                    self.failed_requests[url]["success"] += 1
                    with open(filename, "wb") as file:
                        file.write(response.content)
                    self.chunk_status[chunk_id] = 2
                elif response.status_code >= 400:
                    self.failed_requests[url]["fail"] += 1

                if (
                    self.failed_requests[url]["fail"]
                    / (
                        self.failed_requests[url]["success"]
                        + self.failed_requests[url]["fail"]
                    )
                    > 0.5
                    and self.failed_requests[url]["fail"] > 10
                ):
                    # 如果某 URL 的失败率高于 50%，则跳过该 URL
                    return
            except Exception as e:
                if self.chunk_status[chunk_id] == 1:
                    self.chunk_status[chunk_id] = 0

    def download_file(self, total_size, file_path):
        num_chunks = (total_size + self.chunksize - 1) // self.chunksize
        self.chunk_status = [0] * num_chunks
        try:
            shutil.rmtree(f"temp/{self.listhash}")
        except:
            pass
        os.makedirs(f"temp/{self.listhash}", exist_ok=True)
        max_retry = 3
        while 0 in self.chunk_status:
            if max_retry == 0:
                print("下载出错。")
                return
            with ThreadPoolExecutor(
                max_workers=self.max_conn * len(self.urlist)
            ) as executor:
                for url in self.urlist:
                    for chunk_id in range(num_chunks):
                        executor.submit(self.download_chunk, url, chunk_id, total_size)
            max_retry -= 1

        # 合并所有临时文件到一个文件
        with open(file_path, "wb") as outfile:
            for chunk_id in range(num_chunks):
                filename = f"temp/{self.listhash}/{chunk_id}"
                with open(filename, "rb") as infile:
                    shutil.copyfileobj(infile, outfile)

        # 删除临时目录
        shutil.rmtree(f"temp/{self.listhash}")

        # 验证下载文件
        if os.path.getsize(file_path) != total_size:
            print("文件大小不一致，下载可能出错。")

        return True


def file_download(download_url_list, download_path, request_proxies=None):
    chunksize = 4 * 1024 * 1024  # 分片大小4MB
    max_conn = 4  # 最大连接数
    # 创建对象
    downloader = Downloader(
        download_url_list, chunksize, max_conn, use_proxies=request_proxies
    )

    # 下载文件
    total_size = length(download_url_list)
    if not total_size:
        print("无法连接所有的下载服务器，更新失败")
        return
    print("文件大小已获取，开始下载")
    return downloader.download_file(total_size, download_path)


if __name__ == "__main__":
    while True:
        json_string = input("请输入下载指令：")
        # 解析JSON字符串
        """
        Example:
        {
            "urlist":["https://mirror1.example.com/file.zip","https://mirror2.example.com/file.zip","https://mirror3.example.com/file.zip"],
            "filepath":"C:\\Download\\File.zip"
            "proxies":"http://127.0.0.1:7890"
        }
        """
        try:
            data = json.loads(json_string)
            if isinstance(data, dict):
                # 这两个参数是必须的
                urlist = data.get("urlist", None)
                filepath = data.get("filepath", None)
                # 这个参数是可选的
                proxies = data.get("proxies", None)

                # 对于变量的类型进行检查
                if not isinstance(urlist, list) or not all(
                    isinstance(url, str) for url in urlist
                ):
                    print("urlist类型错误")
                    continue
                if not isinstance(filepath, str):
                    print("filepath类型错误")
                    continue
                if proxies is not None and not isinstance(proxies, str):
                    print("proxies类型错误")
                    continue

                # 启动下载
                file_download(
                    download_url_list=urlist,
                    download_path=filepath,
                    request_proxies=proxies,
                )
            else:
                print("Json结构错误")
        except json.JSONDecodeError:
            print("无法解析JSON。")