from __future__ import annotations
# 用于显示进度条
# from tqdm import tqdm
from tqdm.notebook import tqdm as tqdm
# 用于发起网络请求
import requests

# import signal
#
#
# signal.signal(signal.SIGINT, multitasking.killall)

# 请求头
headers = {
 'User-Agent': 'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.132 Safari/537.36 QIHU 360SE'
}
# 定义 1 MB 多少为 B
MB = 1024 ** 2


def split(start: int, end: int, step: int) -> list[tuple[int, int]]:
 # 分多块
 parts = [(start, min(start + step, end))
 for start in range(0, end, step)]
 return parts


def get_file_size(url: str, raise_error: bool = False) -> int:
 '''
 获取文件大小

 Parameters
 ----------
 url : 文件直链
 raise_error : 如果无法获取文件大小,是否引发错误

 Return
 ------
 文件大小(B为单位)
 如果不支持则会报错

 '''
 response = requests.head(url)
 file_size = response.headers.get('Content-Length')
 if file_size is None:
 if raise_error is True:
 raise ValueError('该文件不支持多线程分段下载!')
 return file_size
 return int(file_size)


def download_multitask(url: str, file_name: str, retry_times: int = 3, each_size=16 * MB) -> None:
 # 用于多线程操作
 import multitasking
 # 导入 retry 库以方便进行下载出错重试
 from retry import retry
 '''
 根据文件直链和文件名下载文件

 Parameters
 ----------
 url : 文件直链
 file_name : 文件名
 retry_times: 可选的,每次连接失败重试次数
 Return
 ------
 None

 '''
 f = open(file_name, 'wb')
 file_size = get_file_size(url)

 @retry(tries=retry_times)
 @multitasking.task
 def start_download(start: int, end: int) -> None:
 '''
 根据文件起止位置下载文件

 Parameters
 ----------
 start : 开始位置
 end : 结束位置
 '''
 _headers = headers.copy()
 # 分段下载的核心
 _headers['Range'] = f'bytes={start}-{end}'
 # 发起请求并获取响应(流式)
 response = session.get(url, headers=_headers, stream=True)
 # 每次读取的流式响应大小
 chunk_size = 128
 # 暂存已获取的响应,后续循环写入
 chunks = []
 for chunk in response.iter_content(chunk_size=chunk_size):
 # 暂存获取的响应
 chunks.append(chunk)
 # 更新进度条
 bar.update(chunk_size)
 f.seek(start)
 for chunk in chunks:
 f.write(chunk)
 # 释放已写入的资源
 del chunks

 session = requests.Session()
 # 分块文件如果比文件大,就取文件大小为分块大小
 each_size = min(each_size, file_size)

 # 分块
 parts = split(0, file_size, each_size)
 print(f'分块数:{len(parts)}')
 # 创建进度条
 bar = tqdm(total=file_size, desc=f'下载文件:{file_name}')
 for part in parts:
 start, end = part
 start_download(start, end)
 # 等待全部线程结束
 multitasking.wait_for_tasks()
 f.close()
 bar.close()


def download_singletask(url, file_name):
 # 导入requests 库
 # import requests
 # 导入 tqdm
 # from tqdm import tqdm, tqdm_notebook
 # from tqdm.notebook import tqdm as tqdm

 # 文件下载直链
 # url = 'https://ffmpeg.org/releases/ffmpeg-6.0.tar.xz'
 # 请求头
 headers = {
 'User-Agent': 'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.132 Safari/537.36 QIHU 360SE'
 }
 # # file_name = 'ffmpeg-6.0.tar.xz'
 # # 发起 head 请求,即只会获取响应头部信息
 # head = requests.head(url, headers=headers)
 # # 文件大小,以 B 为单位
 # file_size = head.headers.get('Content-Length')
 # if file_size is not None:
 # file_size = int(file_size)
 response = requests.get(url, headers=headers, stream=True)
 # 一块文件的大小
 chunk_size = 1024
 file_size = response.headers['content-length']
 if file_size is not None:
 file_size = int(file_size)
 print(file_size)
 bar = tqdm(total=file_size, desc=f'下载文件 {file_name}')
 with open(file_name, mode='wb') as f:
 # 写入分块文件
 for chunk in response.iter_content(chunk_size=chunk_size):
 f.write(chunk)
 bar.update(chunk_size)
 # 关闭进度条
 bar.close()


def download(url, file_name):
 # try:
 # # 用于多线程操作
 # import multitasking
 # # 导入 retry 库以方便进行下载出错重试
 # from retry import retry
 #
 # download_multitask(url, file_name)
 # except Exception as e:
 download_singletask(url, file_name)