from flask import Flask, request, jsonify, Response
import random
import string
import threading
import time
from datetime import datetime, timedelta

app = Flask(__name__)

transfers = {}
transfer_data = {}
transfer_lock = threading.Lock()

# Настройки хранения данных
TRANSFER_LIFETIME = timedelta(hours=1)  # Время хранения завершенных передач
CLEANUP_INTERVAL = 3600  # 5 минут в секундах

def cleanup_task():
    """Фоновая задача для очистки старых передач"""
    while True:
        time.sleep(CLEANUP_INTERVAL)
        with transfer_lock:
            now = datetime.now()
            to_delete = []
            for transfer_id, info in transfers.items():
                if info['completed'] and (now - info['completed_time']) > TRANSFER_LIFETIME:
                    to_delete.append(transfer_id)
            
            for transfer_id in to_delete:
                del transfers[transfer_id]
                del transfer_data[transfer_id]

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

def generate_short_id(length=4):
    while True:
        token = ''.join(random.choices(string.ascii_letters + string.digits, k=length))
        if token not in transfers:
            return token

@app.route('/create_transfer', methods=['POST'])
def create_transfer():
    data = request.json
    transfer_id = generate_short_id()
    
    with transfer_lock:
        transfers[transfer_id] = {
            'filename': data['filename'],
            'filesize': data.get('filesize', 0),
            'completed': False,
            'created_time': datetime.now(),
            'completed_time': None
        }
        transfer_data[transfer_id] = []
    
    return jsonify({'transfer_id': transfer_id})

@app.route('/upload/<transfer_id>', methods=['POST'])
def upload_chunk(transfer_id):
    if transfer_id not in transfers:
        return jsonify({'error': 'Invalid transfer ID'}), 404

    chunk = request.data
    with transfer_lock:
        transfer_data[transfer_id].append(chunk)
        
        if request.headers.get('X-Transfer-Complete') == 'true':
            transfers[transfer_id]['completed'] = True
            transfers[transfer_id]['completed_time'] = datetime.now()
    
    return jsonify({'status': 'chunk accepted'})

@app.route('/download/<transfer_id>', methods=['GET'])
def download_file(transfer_id):
    def generate():
        index = 0
        while True:
            with transfer_lock:
                transfer = transfers.get(transfer_id)
                if not transfer:
                    break
                
                chunks = transfer_data.get(transfer_id, [])
                
                # Отправляем все новые чанки
                while index < len(chunks):
                    yield chunks[index]
                    index += 1
                
                # Проверяем завершение передачи
                if transfer['completed']:
                    break
            
            time.sleep(1)  # Ожидание новых чанков

    if transfer_id not in transfers:
        return jsonify({'error': 'Transfer not found'}), 404

    return Response(
        generate(),
        mimetype='application/octet-stream',
        headers={
            'Content-Disposition': f'attachment; filename="{transfers[transfer_id]["filename"]}"',
            'Transfer-Encoding': 'chunked'
        }
    )

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