from flask import Flask, request, jsonify, Response
import random
import string
import time
import threading
import os
from collections import defaultdict

app = Flask(__name__)

# Конфигурация
MAX_MEMORY = 12 * 1024 * 1024 * 1024  # 12 GB
CHUNK_SIZE = 5 * 1024 * 1024  # 5 MB
FILE_STORAGE = "./temp_uploads"
os.makedirs(FILE_STORAGE, exist_ok=True)

active_transfers = defaultdict(dict)
transfer_lock = threading.Lock()

def generate_transfer_id():
    return ''.join(random.choices(string.ascii_letters + string.digits, k=16))

@app.route('/init_upload', methods=['POST'])
def init_upload():
    filename = request.json.get('filename')
    filesize = request.json.get('filesize')
    
    transfer_id = generate_transfer_id()
    filepath = os.path.join(FILE_STORAGE, transfer_id)
    
    with transfer_lock:
        active_transfers[transfer_id] = {
            'filename': filename,
            'filesize': filesize,
            'received': 0,
            'filepath': filepath,
            'created_at': time.time()
        }
    
    # Создаем пустой файл
    with open(filepath, 'wb') as f:
        f.truncate(filesize)
    
    return jsonify({'transfer_id': transfer_id})

@app.route('/upload/<transfer_id>', methods=['POST'])
def upload_chunk(transfer_id):
    if transfer_id not in active_transfers:
        return jsonify({'error': 'Invalid transfer ID'}), 404
    
    chunk = request.data
    start = int(request.headers.get('Content-Range').split('-')[0])
    
    filepath = active_transfers[transfer_id]['filepath']
    
    try:
        with open(filepath, 'r+b') as f:
            f.seek(start)
            f.write(chunk)
        
        with transfer_lock:
            active_transfers[transfer_id]['received'] += len(chunk)
            
        return jsonify({'status': 'success', 'received': active_transfers[transfer_id]['received']})
    
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/download/<transfer_id>', methods=['GET'])
def download_file(transfer_id):
    def generate():
        filepath = active_transfers[transfer_id]['filepath']
        with open(filepath, 'rb') as f:
            while True:
                chunk = f.read(CHUNK_SIZE)
                if not chunk:
                    break
                yield chunk
        
        # Удаляем файл после отправки
        os.remove(filepath)
        del active_transfers[transfer_id]
    
    return Response(
        generate(),
        mimetype='application/octet-stream',
        headers={
            'Content-Disposition': f'attachment; filename="{active_transfers[transfer_id]["filename"]}"',
            'Content-Length': str(active_transfers[transfer_id]['filesize'])
        }
    )

# Очистка зависших передач
def cleanup():
    while True:
        time.sleep(3600)
        now = time.time()
        for tid in list(active_transfers.keys()):
            if now - active_transfers[tid]['created_at'] > 86400:  # 24 часа
                try:
                    os.remove(active_transfers[tid]['filepath'])
                    del active_transfers[tid]
                except:
                    pass

threading.Thread(target=cleanup, daemon=True).start()

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, threaded=True)