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

app = Flask(__name__)

transfers = {}
transfer_data = {}
transfer_lock = threading.Lock()
CLEANUP_INTERVAL = 300  # 5 минут

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

def cleanup_old_transfers():
    while True:
        now = datetime.now()
        with transfer_lock:
            to_delete = []
            for transfer_id, data in transfers.items():
                if now - data['last_activity'] > timedelta(seconds=CLEANUP_INTERVAL):
                    to_delete.append(transfer_id)
            
            for transfer_id in to_delete:
                del transfers[transfer_id]
                del transfer_data[transfer_id]
        
        threading.Event().wait(CLEANUP_INTERVAL)

@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(8)
    with transfer_lock:
        transfers[transfer_id] = {
            'filename': filename,
            'filesize': filesize,
            'created': datetime.now(),
            'last_activity': datetime.now(),
            '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):
    with transfer_lock:
        if transfer_id not in transfers:
            return jsonify({'error': 'Invalid transfer ID'}), 404
        
        transfers[transfer_id]['last_activity'] = datetime.now()
        
        data = request.data
        if data:
            transfer_data[transfer_id].append(data)
        
        if request.headers.get('X-Transfer-Complete') == 'true':
            transfers[transfer_id]['completed'] = True
        
        return jsonify({'status': 'OK', 'received': len(data)})

@app.route('/download/<transfer_id>', methods=['GET'])
def download_file(transfer_id):
    def generate():
        with transfer_lock:
            if transfer_id not in transfers:
                return
            
            transfers[transfer_id]['last_activity'] = datetime.now()
            transfer = transfers[transfer_id]
            data_queue = transfer_data[transfer_id]
        
        while True:
            with transfer_lock:
                if data_queue:
                    chunk = bytes(data_queue.popleft())
                    yield chunk
                elif transfer['completed']:
                    break
                else:
                    # Ждем новые данные
                    pass
            
            threading.Event().wait(0.1)

    if transfer_id not in transfers:
        return jsonify({'error': 'Invalid transfer ID'}), 404
    
    return Response(
        generate(),
        mimetype='application/octet-stream',
        headers={
            'Content-Disposition': f'attachment; filename="{transfers[transfer_id]["filename"]}"',
            'Content-Length': str(transfers[transfer_id]['filesize'])
        }
    )

if __name__ == '__main__':
    threading.Thread(target=cleanup_old_transfers, daemon=True).start()
    app.run(host='0.0.0.0', port=7860, threaded=True)