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

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'temp_uploads'
Path(app.config['UPLOAD_FOLDER']).mkdir(exist_ok=True)

# Конфигурация памяти
MAX_MEMORY = 12 * 1024 * 1024 * 1024  # 12 GB
CHUNK_SIZE = 100 * 1024 * 1024  # 4 MB
MAX_FILE_SIZE = 20 * 1024 * 1024 * 1024  # 20 GB

transfers = {}
active_memory = 0
lock = threading.Lock()

class TransferManager:
    def __init__(self):
        self.transfers = {}
        self.file_locks = {}
    
    def create_transfer(self, filename, filesize):
        if filesize > MAX_FILE_SIZE:
            raise ValueError("File size exceeds maximum limit")
            
        transfer_id = ''.join(random.choices(string.ascii_letters + string.digits, k=16))
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], transfer_id)
        
        with lock:
            self.transfers[transfer_id] = {
                'filename': filename,
                'filesize': filesize,
                'received': 0,
                'file_path': file_path,
                'created_at': time.time(),
                'status': 'uploading'
            }
            self.file_locks[transfer_id] = threading.Lock()
        
        return transfer_id

manager = TransferManager()

@app.route('/init_upload', methods=['POST'])
def init_upload():
    filename = request.json.get('filename')
    filesize = request.json.get('filesize')
    
    try:
        transfer_id = manager.create_transfer(filename, filesize)
        return jsonify({'transfer_id': transfer_id})
    except Exception as e:
        return jsonify({'error': str(e)}), 400

@app.route('/upload/<transfer_id>', methods=['POST'])
def upload_chunk(transfer_id):
    if transfer_id not in manager.transfers:
        return jsonify({'error': 'Invalid transfer ID'}), 404
    
    chunk = request.data
    chunk_size = len(chunk)
    
    with lock:
        transfer = manager.transfers[transfer_id]
        if transfer['status'] != 'uploading':
            return jsonify({'error': 'Transfer completed'}), 400
    
    try:
        with manager.file_locks[transfer_id], open(transfer['file_path'], 'ab') as f:
            f.write(chunk)
            transfer['received'] += chunk_size
        
        if transfer['received'] >= transfer['filesize']:
            transfer['status'] = 'ready'
        
        return jsonify({
            'received': transfer['received'],
            'status': transfer['status']
        })
    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 = manager.transfers.get(transfer_id, {}).get('file_path')
        if not file_path or not os.path.exists(file_path):
            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
        
        # Автоочистка через 1 час
        threading.Timer(3600, cleanup_transfer, args=[transfer_id]).start()
    
    return Response(
        generate(),
        mimetype='application/octet-stream',
        headers={'Content-Disposition': f'attachment; filename="{manager.transfers[transfer_id]["filename"]}"'}
    )

def cleanup_transfer(transfer_id):
    with lock:
        if transfer_id in manager.transfers:
            try:
                os.remove(manager.transfers[transfer_id]['file_path'])
            except:
                pass
            del manager.transfers[transfer_id]

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