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

app = Flask(__name__)

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

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():
    filename = request.json.get('filename')
    filesize = request.json.get('filesize', 0)
    
    transfer_id = generate_short_id(4)
    with transfer_lock:
        transfers[transfer_id] = {
            'filename': filename,
            'filesize': filesize,
            'completed': False
        }
        transfer_data[transfer_id] = deque()
    return jsonify({'transfer_id': transfer_id})

@app.route('/upload/<transfer_id>', methods=['POST'])
def upload_file(transfer_id):
    if transfer_id not in transfers:
        return jsonify({'error': 'Invalid transfer ID'}), 404
    
    data = request.data
    with transfer_lock:
        transfer_data[transfer_id].append(data)
        transfers[transfer_id]['completed'] = request.headers.get('X-Transfer-Complete') == 'true'
    
    return jsonify({'status': 'chunk uploaded'})

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

    def generate():
        while True:
            with transfer_lock:
                if transfer_data[transfer_id]:
                    chunk = transfer_data[transfer_id].popleft()
                    yield chunk
                elif transfers[transfer_id]['completed']:
                    break

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

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