from flask import Flask, request, Response
import random
import string
from collections import deque
import threading
import time

app = Flask(__name__)

transfers = {}
transfer_lock = threading.Lock()
BUFFER_SIZE = 5  # Максимальное количество хранимых чанков

class TransferSession:
    def __init__(self, filename):
        self.filename = filename
        self.buffer = deque(maxlen=BUFFER_SIZE)
        self.completed = False
        self.event = threading.Event()
        self.last_activity = time.time()

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

@app.route('/start_upload', methods=['POST'])
def start_upload():
    filename = request.json.get('filename', 'file.bin')
    transfer_id = generate_id()
    
    with transfer_lock:
        transfers[transfer_id] = TransferSession(filename)
    
    return {'transfer_id': transfer_id}

@app.route('/upload/<transfer_id>', methods=['POST'])
def upload_chunk(transfer_id):
    with transfer_lock:
        session = transfers.get(transfer_id)
        if not session:
            return {'error': 'Invalid transfer ID'}, 404
        
        chunk = request.data
        session.buffer.append(chunk)
        session.last_activity = time.time()
        session.event.set()
        session.event.clear()
        
        if request.headers.get('X-Transfer-Complete') == 'true':
            session.completed = True
        
        return {'status': 'chunk accepted'}

@app.route('/download/<transfer_id>')
def download_chunks(transfer_id):
    def generate():
        with transfer_lock:
            session = transfers.get(transfer_id)
            if not session:
                yield b'Transfer not found'
                return
            
            while True:
                while session.buffer:
                    chunk = session.buffer.popleft()
                    yield chunk
                
                if session.completed:
                    break
                
                session.event.wait(timeout=30)
                if time.time() - session.last_activity > 60:
                    break

        # Очистка через минуту после завершения
        time.sleep(60)
        with transfer_lock:
            if transfer_id in transfers:
                del transfers[transfer_id]

    return Response(
        generate(),
        mimetype='application/octet-stream',
        headers={'Content-Disposition': f'attachment; filename="{session.filename}"'}
    )

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