from flask import Flask, send_from_directory, request, jsonify
import os
import threading
from concurrent.futures import ThreadPoolExecutor
import queue
import time

"""
增加最大下载数量，可作为备用服务
"""

app = Flask(__name__)

# 设置允许访问的文件夹路径
UPLOAD_FOLDER = r'C:\小工具'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# 创建一个队列来管理下载请求
download_queue = queue.Queue()

# 创建一个线程池，最大工作线程数为5
executor = ThreadPoolExecutor(max_workers=5)

def calculate_md5(file_path):
    """计算文件的 MD5 哈希值"""
    hash_md5 = hashlib.md5()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()

@app.route('/files', methods=['GET'])
def list_files():
    # 获取指定目录下的所有文件名
    files = os.listdir(app.config['UPLOAD_FOLDER'])
    return jsonify(files)

@app.route('/download/<filename>', methods=['GET'])
def download_file(filename):
    # 检查文件是否存在
    file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    if not os.path.exists(file_path):
        return "File not found", 404

    # 将下载请求加入队列
    download_queue.put((filename, request.environ['REMOTE_ADDR']))

    # 返回一个消息，告知客户端请求已加入队列
    return jsonify({"message": "Download request added to queue"})

@app.route('/hash/<filename>', methods=['GET'])
def get_file_hash(filename):
    """返回文件的 MD5 哈希值"""
    file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    if os.path.exists(file_path):
        md5_hash = calculate_md5(file_path)
        return md5_hash
    else:
        return "File not found", 404

def process_download_queue():
    while True:
        # 从队列中获取下一个下载请求
        filename, client_ip = download_queue.get()

        # 使用线程池提交下载任务
        future = executor.submit(download_and_release, filename, client_ip)

        # 等待任务完成
        future.result()

def download_and_release(filename, client_ip):
    file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    try:
        # 发送文件给客户端
        with app.test_request_context():
            response = send_from_directory(app.config['UPLOAD_FOLDER'], filename, as_attachment=True)
            response.direct_passthrough = False  # 确保整个文件一次性发送
            response.headers['X-Client-IP'] = client_ip  # 添加客户端IP头，以便后续处理
            return response
    except Exception as e:
        return str(e), 404

if __name__ == '__main__':
    # 确保文件夹存在
    if not os.path.exists(app.config['UPLOAD_FOLDER']):
        os.makedirs(app.config['UPLOAD_FOLDER'])

    # 启动一个单独的线程来处理下载队列
    threading.Thread(target=process_download_queue, daemon=True).start()

    app.run(host='192.168.10.86', debug=True, port=5001)