import cProfile
import pstats
import sys

from aiohttp import web

profiler = None
import io


@web.middleware
async def profile_middleware(request, handler):
    global profiler
    if not profiler:
        profiler = cProfile.Profile()
    profiler.enable()
    response = await handler(request)
    profiler.disable()
    stats = pstats.Stats(profiler, stream=sys.stdout)
    stats.sort_stats("cumulative")
    stats.print_stats()
    return response


async def profiler_handler(request):
    output = io.StringIO()
    stats = pstats.Stats(profiler, stream=output)
    sort_by = request.query.get("sort", "tot. percall")
    stats.sort_stats(sort_by)
    stats.print_stats()
    return web.Response(text=output.getvalue())


class Profiler:

    def __init__(self):
        global profiler
        if profiler is None:
            profiler = cProfile.Profile()
        self.profiler = profiler

    async def __aenter__(self):
        self.profiler.enable()

    async def __aexit__(self, *args, **kwargs):
        self.profiler.disable()
