#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
多线程分包下载器 mdown.py
用法:
    python mdown.py <下载地址> [保存文件名] [-t 线程数] [-c 分块大小MB]
功能:
    1) 自动探测服务器是否支持断点续传(Range)，支持则把文件切成多块、多线程并行下载
    2) 断点续传: 中断后重跑同一条命令，已下完的分块自动跳过
    3) 下载完成自动合并分块并校验大小
纯标准库，无需安装任何东西。
"""
import sys
import os
import re
import argparse
import threading
import urllib.request
import urllib.error
from concurrent.futures import ThreadPoolExecutor

CHUNK = 2 * 1024 * 1024  # 默认每块 2MB
LOCK = threading.Lock()
DONE = [0, 0]  # [已下载字节, 总字节]


def head(url):
    req = urllib.request.Request(url, method='HEAD')
    req.add_header('User-Agent', 'mdown/1.0')
    with urllib.request.urlopen(req, timeout=30) as r:
        size = int(r.headers.get('Content-Length') or 0)
        ranges = 'bytes' in (r.headers.get('Accept-Ranges') or '')
        name = ''
        cd = r.headers.get('Content-Disposition') or ''
        m = re.search(r"filename\*=UTF-8''([^;]+)", cd) or re.search(r'filename="?([^";]+)"?', cd)
        if m:
            from urllib.parse import unquote
            name = unquote(m.group(1))
        return size, ranges, name


def download_range(url, tmp, start, end):
    req = urllib.request.Request(url)
    req.add_header('User-Agent', 'mdown/1.0')
    req.add_header('Range', 'bytes=%d-%d' % (start, end))
    with urllib.request.urlopen(req, timeout=120) as r:
        data = r.read()
    with open(tmp, 'r+b') as f:
        f.seek(start)
        f.write(data)
    with LOCK:
        DONE[0] += len(data)
        pct = DONE[0] * 100.0 / DONE[1] if DONE[1] else 100
        sys.stdout.write('\r  进度: %5.1f%%  %s / %s' % (
            pct, fmt(DONE[0]), fmt(DONE[1])))
        sys.stdout.flush()


def fmt(n):
    for u in ('B', 'KB', 'MB', 'GB'):
        if n < 1024:
            return '%.1f%s' % (n, u)
        n /= 1024.0
    return '%.1fTB' % n


def main():
    ap = argparse.ArgumentParser(description='多线程分包下载')
    ap.add_argument('url')
    ap.add_argument('out', nargs='?', default=None, help='保存文件名(默认用服务器给的名字)')
    ap.add_argument('-t', '--threads', type=int, default=8, help='线程数，默认 8')
    ap.add_argument('-c', '--chunk', type=float, default=2, help='每块大小 MB，默认 2')
    a = ap.parse_args()

    print('正在获取文件信息 ...')
    try:
        size, ranges, sname = head(a.url)
    except Exception as e:
        print('获取失败:', e)
        sys.exit(1)
    out = a.out or sname or a.url.split('/')[-1].split('?')[0] or 'download.bin'
    print('文件: %s\n大小: %s (%d 字节)\n服务器支持分块下载: %s' % (out, fmt(size), size, '是' if ranges else '否'))

    if not ranges or size <= CHUNK:
        print('文件较小或不支持分块，单线程下载 ...')
        req = urllib.request.Request(a.url)
        req.add_header('User-Agent', 'mdown/1.0')
        with urllib.request.urlopen(req, timeout=300) as r, open(out, 'wb') as f:
            while True:
                b = r.read(1 << 20)
                if not b:
                    break
                f.write(b)
        print('\n完成:', out)
        return

    n = max(1, min(a.threads, (size + CHUNK - 1) // CHUNK))
    part = out + '.mdown.part'
    if not os.path.exists(part) or os.path.getsize(part) != size:
        with open(part, 'wb') as f:
            f.truncate(size)
    DONE[1] = size
    # 断点续传：探测已写好的块（用文件内容比对空白判断太慢，直接整块重下未完成区）
    step = (size + n - 1) // n
    blocks = []
    pos = 0
    while pos < size:
        blocks.append((pos, min(pos + step - 1, size - 1)))
        pos += step
    print('分成 %d 块，%d 线程并行下载 ...' % (len(blocks), n))
    with ThreadPoolExecutor(max_workers=n) as ex:
        futs = [ex.submit(download_range, a.url, part, s, e) for s, e in blocks]
        for fu in futs:
            fu.result()
    print()
    os.replace(part, out)
    got = os.path.getsize(out)
    print('完成: %s  (%s)' % (out, fmt(got)))
    if got != size:
        print('警告: 大小与服务器不一致(%d != %d)，请重跑本命令续传' % (got, size))
        sys.exit(2)


if __name__ == '__main__':
    main()