from flask import Flask, request, jsonify, Response
import random
import string
import threading
import time
import os
from pathlib import Path

app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 12 * 1024 * 1024 * 1024  # 12GB
UPLOAD_FOLDER = Path('uploads')
UPLOAD_FOLDER.mkdir(exist_ok=True)

transfers = {}
transfer_lock = threading.Lock()
TIMEOUT = 3600  # 1 час
CHUNK_SIZE = 1024 * 1024  # 1MB

def generate_short_id(length=10):
    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()
    file_path = UPLOAD_FOLDER / f"{transfer_id}.dat"
    
    with transfer_lock:
        transfers[transfer_id] = {
            'filename': filename,
            'filesize': filesize,
            'completed': False,
            'last_activity': time.time(),
            'file_path': file_path,
            'file_handle': open(file_path, 'wb')
        }
        
    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
        
        transfer = transfers[transfer_id]
        transfer['last_activity'] = time.time()
        
        try:
            # Записываем полученные данные в файл
            data = request.data
            transfer['file_handle'].write(data)
            transfer['file_handle'].flush()
            
            if request.headers.get('X-Transfer-Complete') == 'true':
                transfer['completed'] = True
                transfer['file_handle'].close()
                
            return jsonify({'status': 'chunk uploaded', 'received': len(data)})
        
        except Exception as e:
            return jsonify({'error': str(e)}), 500

@app.route('/download/<transfer_id>', methods=['GET'])
def download_file(transfer_id):
    def generate():
        file_path = None
        with transfer_lock:
            if transfer_id not in transfers:
                yield b'Transfer expired'
                return
                
            transfer = transfers[transfer_id]
            file_path = transfer['file_path']
            transfer['last_activity'] = time.time()
        
        if not file_path.exists():
            yield b'File not found'
            return
            
        with open(file_path, 'rb') as f:
            while True:
                chunk = f.read(CHUNK_SIZE)
                if not chunk:
                    break
                yield chunk

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

def cleanup_thread():
    while True:
        time.sleep(60 * 5)  # Проверка каждые 5 минут
        with transfer_lock:
            now = time.time()
            expired = []
            for tid, t in transfers.items():
                if now - t['last_activity'] > TIMEOUT:
                    expired.append(tid)
                    if 'file_handle' in t and not t['file_handle'].closed:
                        t['file_handle'].close()
                    if t['file_path'].exists():
                        t['file_path'].unlink()
            
            for tid in expired:
                del transfers[tid]

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

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