import tempfile

from flask import current_app
import werkzeug.formparser
import werkzeug.wrappers

def werkzeug_patch():
    global werkzeug_orig_stream_factory

    werkzeug_orig_stream_factory = werkzeug.formparser.default_stream_factory

    def custom_stream_factory(total_content_length, filename, content_type, content_length=None):
        if total_content_length > 1024 * 500:
                return tempfile.NamedTemporaryFile('wb+', prefix = 'upload_', dir = current_app.config['FILE_DIRECTORY'], delete = True)
        return werkzeug_orig_stream_factory(total_content_length, filename, content_type, content_length)

    werkzeug.formparser.default_stream_factory = custom_stream_factory
    werkzeug.wrappers.default_stream_factory = custom_stream_factory

def werkzeug_reset():
    werkzeug.formparser.default_stream_factory = werkzeug_orig_stream_factory
    werkzeug.wrappers.default_stream_factory = werkzeug_orig_stream_factory

def patch():
    werkzeug_patch()