"""Using pyinstrument to profile FastHTML apps."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['ProfileMiddleware', 'instrument']

# %% ../nbs/00_core.ipynb 3
from fasthtml.common import *
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.middleware import Middleware
from functools import wraps
try: from rich import print
except: pass
try:
    from pyinstrument import Profiler
except ImportError:
    raise ImportError('Please install pyinstrument')

# %% ../nbs/00_core.ipynb 4
class ProfileMiddleware(BaseHTTPMiddleware):    
    async def dispatch(self, request, call_next):
        profiling = request.query_params.get("profile", False)
        term = request.query_params.get("term", False)  
        if profiling:
            profiler = Profiler()
            profiler.start()
            response = await call_next(request)            
            profiler.stop()
            if term: 
                print(profiler.output_text())
                return response
            return HTMLResponse(profiler.output_html())
        return await call_next(request)

# %% ../nbs/00_core.ipynb 5
def instrument(route_handler):
    @functools.wraps(route_handler)
    def _wrapper(*args, **kwargs):
        # Profile the action
        profiler = Profiler()
        profiler.start()
        route_handler(*args, **kwargs)
        profiler.stop()
        # Return the profile output
        return profiler.output_html()
    return _wrapper
